Embedding quant unification (pytorch#14622)

Summary: Pull Request resolved: pytorch#14622

Differential Revision: D83318725

Author: metascroy

.ci/scripts/test_llama.sh

@@ -237,7 +237,7 @@ if [[ "${CUSTOM}" == "ON" ]]; then
   EXPORT_ARGS="${EXPORT_ARGS} model.use_sdpa_with_kv_cache=true"
 fi
 if [[ "${QE}" == "ON" ]]; then
-  EXPORT_ARGS="${EXPORT_ARGS} quantization.embedding_quantize=\"8,1024\""
+  EXPORT_ARGS="${EXPORT_ARGS} quantization.embedding_quantize=\"8,768\""
 fi
 if [[ "${MPS}" == "ON" ]]; then
   EXPORT_ARGS="${EXPORT_ARGS} backend.mps.enabled=true model.enable_dynamic_shape=false debug.verbose=true"

examples/models/llama/source_transformation/quantize.py

@@ -595,19 +595,12 @@ def __init__(
 
     @torch.no_grad()
     def create_quantized_state_dict(self, packed=False) -> Dict:
+        from torchao.quantization.quant_api import IntxWeightOnlyConfig, quantize_, MappingType
+        from torchao.quantization.granularity import PerGroup
+
         cur_state_dict = self.mod.state_dict()
 
-        if self.bitwidth == 2:
-            range_min = -2
-            range_max = 1
-        elif self.bitwidth == 4:
-            range_min = -8
-            range_max = 7
-        elif self.bitwidth == 8:
-            range_min = -128
-            range_max = 127
-        else:
-            raise ValueError(f"Unsupported bitwidth {self.bitwidth}")
+        assert self.bitwidth in [2, 4, 8], f"Unsupported bitwidth {self.bitwidth}"
 
         for fqn, mod in self.mod.named_modules():
             if isinstance(mod, nn.Embedding):
@@ -619,18 +612,18 @@ def create_quantized_state_dict(self, packed=False) -> Dict:
                 print(
                     f"quantize {fqn, mod} with group_size {self.group_size}, bitwidth {self.bitwidth}"
                 )
-                weight, scales, _ = dynamically_quantize_per_channel(
-                    (
-                        mod.weight.to(dtype=self.precision)
-                        if self.precision
-                        else mod.weight
-                    ),
-                    range_min,
-                    range_max,
-                    torch.int8,
-                    self.group_size,
-                    scales_dtype=mod.weight.dtype,
+                tmp_model = nn.Embedding(mod.weight.shape[0], mod.weight.shape[1])
+                if self.precision:
+                    tmp_model = tmp_model.to(dtype=self.precision)
+                tmp_model.weight = nn.Parameter(mod.weight)
+                config = IntxWeightOnlyConfig(
+                    weight_dtype=getattr(torch, f"int{self.bitwidth}"),
+                    granularity=PerGroup(self.group_size),
+                    mapping_type=MappingType.SYMMETRIC,
                 )
+                quantize_(tmp_model, config, lambda m, fqn: isinstance(m, nn.Embedding))
+                weight = tmp_model.weight.qdata
+                scales = tmp_model.weight.scale
 
                 if packed:
                     if self.bitwidth == 2:

exir/passes/_quant_patterns_and_replacements.py

@@ -986,25 +986,46 @@ def replacement(x, dim, start, end, x_scale, x_zero_point, x_qmin, x_qmax):
     ]
 
 
-def _get_embedding_ops_patterns_and_replacements_torchao() -> (  # noqa C901
+def _get_embedding_ops_patterns_and_replacements_torchao(node_value_dict) -> (  # noqa C901
     List[Tuple[Callable, Callable, List[Callable]]]
 ):
+    assert node_value_dict is not None, "node_value_dict cannot be None"
+    def get_embedding_replacement_filter(has_nonzero_zero_point):
+        def _filter(match, original_graph, pattern_graph):
+            def get_val(name):
+                node = [n for n in match.nodes_map if n.name == name][0]
+                val = match.nodes_map[node]
+                if isinstance(val, torch.fx.Node) and val.target in node_value_dict:
+                    return node_value_dict[val.target]
+                return val
+
+            zero_point = get_val("zero_point")
+            all_zero = (zero_point == 0).all().item()
+            if has_nonzero_zero_point:
+                return not all_zero
+            else:
+                return all_zero
+        return _filter
+
     def embedding_byte_pattern(indices, int_data, group_size, scale, zero_point):
         dq = torch.ops.torchao.dequantize_affine.default(
             int_data, [1, group_size], scale, zero_point, torch.int8, -128, 127
         )
         return torch.ops.aten.embedding.default(dq, indices)
 
-    def embedding_byte_replacement(indices, int_data, group_size, scale, zero_point):
-        zero_point_dtype_cast = torch.ops.aten.to.dtype(zero_point, scale.dtype)
-        return torch.ops.quantized_decomposed.embedding_byte.default(
-            int_data,
-            scale,
-            zero_point_dtype_cast,
-            -128,
-            127,
-            indices,
-        )
+    def get_embedding_byte_replacement(has_nonzero_zero_point):
+        def embedding_byte_replacement(indices, int_data, group_size, scale, zero_point):
+            zero_point_dtype_cast = torch.ops.aten.to.dtype(zero_point, scale.dtype)
+            zero_point_dtype_cast = zero_point_dtype_cast if has_nonzero_zero_point else None
+            return torch.ops.quantized_decomposed.embedding_byte.default(
+                int_data,
+                scale,
+                zero_point_dtype_cast,
+                -128,
+                127,
+                indices,
+            )
+        return embedding_byte_replacement
 
     def embedding_byte_dtype_pattern(
         indices, int_data, group_size, scale, zero_point, output_dtype
@@ -1021,34 +1042,40 @@ def embedding_byte_dtype_pattern(
         )
         return torch.ops.aten.embedding.default(dq, indices)
 
-    def embedding_byte_dtype_replacement(
-        indices, int_data, group_size, scale, zero_point, output_dtype
-    ):
-        zero_point_dtype_cast = torch.ops.aten.to.dtype(zero_point, scale.dtype)
-        return torch.ops.quantized_decomposed.embedding_byte.dtype(
-            int_data,
-            scale,
-            zero_point_dtype_cast,
-            -128,
-            127,
-            indices,
-            dtype=output_dtype,
-        )
+    def get_embedding_byte_dtype_replacement(has_nonzero_zero_point):
+        def embedding_byte_dtype_replacement(
+            indices, int_data, group_size, scale, zero_point, output_dtype
+        ):
+            zero_point_dtype_cast = torch.ops.aten.to.dtype(zero_point, scale.dtype)
+            zero_point_dtype_cast = zero_point_dtype_cast if has_nonzero_zero_point else None
+            return torch.ops.quantized_decomposed.embedding_byte.dtype(
+                int_data,
+                scale,
+                zero_point_dtype_cast,
+                -128,
+                127,
+                indices,
+                dtype=output_dtype,
+            )
+        return embedding_byte_dtype_replacement
 
     def embedding_2bit_pattern(indices, int_data, group_size, scale, zero_point):
         dq = torch.ops.torchao.dequantize_affine.default(
             int_data, [1, group_size], scale, zero_point, torch.int8, -2, 1
         )
         return torch.ops.aten.embedding.default(dq, indices)
 
-    def embedding_2bit_replacement(indices, int_data, group_size, scale, zero_point):
-        packed_int_data = torch.ops.quant_fusion._pack_embedding_weight.default(
-            int_data, 2
-        )
-        zero_point_dtype_cast = torch.ops.aten.to.dtype(zero_point, scale.dtype)
-        return torch.ops.quantized_decomposed.embedding_2bit.default(
-            packed_int_data, scale, zero_point_dtype_cast, -2, 1, indices
-        )
+    def get_embedding_2bit_replacement(has_nonzero_zero_point):
+        def embedding_2bit_replacement(indices, int_data, group_size, scale, zero_point):
+            packed_int_data = torch.ops.quant_fusion._pack_embedding_weight.default(
+                int_data, 2
+            )
+            zero_point_dtype_cast = torch.ops.aten.to.dtype(zero_point, scale.dtype)
+            zero_point_dtype_cast = zero_point_dtype_cast if has_nonzero_zero_point else None
+            return torch.ops.quantized_decomposed.embedding_2bit.default(
+                packed_int_data, scale, zero_point_dtype_cast, -2, 1, indices
+            )
+        return embedding_2bit_replacement
 
     def embedding_2bit_dtype_pattern(
         indices, int_data, group_size, scale, zero_point, output_dtype
@@ -1065,37 +1092,43 @@ def embedding_2bit_dtype_pattern(
         )
         return torch.ops.aten.embedding.default(dq, indices)
 
-    def embedding_2bit_dtype_replacement(
-        indices, int_data, group_size, scale, zero_point, output_dtype
-    ):
-        packed_int_data = torch.ops.quant_fusion._pack_embedding_weight.default(
-            int_data, 2
-        )
-        zero_point_dtype_cast = torch.ops.aten.to.dtype(zero_point, scale.dtype)
-        return torch.ops.quantized_decomposed.embedding_2bit.dtype(
-            packed_int_data,
-            scale,
-            zero_point_dtype_cast,
-            -2,
-            1,
-            indices,
-            dtype=output_dtype,
-        )
+    def get_embedding_2bit_dtype_replacement(has_nonzero_zero_point):
+        def embedding_2bit_dtype_replacement(
+            indices, int_data, group_size, scale, zero_point, output_dtype
+        ):
+            packed_int_data = torch.ops.quant_fusion._pack_embedding_weight.default(
+                int_data, 2
+            )
+            zero_point_dtype_cast = torch.ops.aten.to.dtype(zero_point, scale.dtype)
+            zero_point_dtype_cast = zero_point_dtype_cast if has_nonzero_zero_point else None
+            return torch.ops.quantized_decomposed.embedding_2bit.dtype(
+                packed_int_data,
+                scale,
+                zero_point_dtype_cast,
+                -2,
+                1,
+                indices,
+                dtype=output_dtype,
+            )
+        return embedding_2bit_dtype_replacement
 
     def embedding_4bit_pattern(indices, int_data, group_size, scale, zero_point):
         dq = torch.ops.torchao.dequantize_affine.default(
             int_data, [1, group_size], scale, zero_point, torch.int8, -8, 7
         )
         return torch.ops.aten.embedding.default(dq, indices)
 
-    def embedding_4bit_replacement(indices, int_data, group_size, scale, zero_point):
-        packed_int_data = torch.ops.quant_fusion._pack_embedding_weight.default(
-            int_data, 4
-        )
-        zero_point_dtype_cast = torch.ops.aten.to.dtype(zero_point, scale.dtype)
-        return torch.ops.quantized_decomposed.embedding_4bit.default(
-            packed_int_data, scale, zero_point_dtype_cast, -8, 7, indices
-        )
+    def get_embedding_4bit_replacement(has_nonzero_zero_point):
+        def embedding_4bit_replacement(indices, int_data, group_size, scale, zero_point):
+            packed_int_data = torch.ops.quant_fusion._pack_embedding_weight.default(
+                int_data, 4
+            )
+            zero_point_dtype_cast = torch.ops.aten.to.dtype(zero_point, scale.dtype)
+            zero_point_dtype_cast = zero_point_dtype_cast if has_nonzero_zero_point else None
+            return torch.ops.quantized_decomposed.embedding_4bit.default(
+                packed_int_data, scale, zero_point_dtype_cast, -8, 7, indices
+            )
+        return embedding_4bit_replacement
 
     def embedding_4bit_dtype_pattern(
         indices, int_data, group_size, scale, zero_point, output_dtype
@@ -1112,53 +1145,86 @@ def embedding_4bit_dtype_pattern(
         )
         return torch.ops.aten.embedding.default(dq, indices)
 
-    def embedding_4bit_dtype_replacement(
-        indices, int_data, group_size, scale, zero_point, output_dtype
-    ):
-        packed_int_data = torch.ops.quant_fusion._pack_embedding_weight.default(
-            int_data, 4
-        )
-        zero_point_dtype_cast = torch.ops.aten.to.dtype(zero_point, scale.dtype)
-        return torch.ops.quantized_decomposed.embedding_4bit.dtype(
-            packed_int_data,
-            scale,
-            zero_point_dtype_cast,
-            -8,
-            7,
-            indices,
-            dtype=output_dtype,
-        )
+    def get_embedding_4bit_dtype_replacement(has_nonzero_zero_point):
+        def embedding_4bit_dtype_replacement(
+            indices, int_data, group_size, scale, zero_point, output_dtype
+        ):
+            packed_int_data = torch.ops.quant_fusion._pack_embedding_weight.default(
+                int_data, 4
+            )
+            zero_point_dtype_cast = torch.ops.aten.to.dtype(zero_point, scale.dtype)
+            zero_point_dtype_cast = zero_point_dtype_cast if has_nonzero_zero_point else None
+            return torch.ops.quantized_decomposed.embedding_4bit.dtype(
+                packed_int_data,
+                scale,
+                zero_point_dtype_cast,
+                -8,
+                7,
+                indices,
+                dtype=output_dtype,
+            )
+        return embedding_4bit_dtype_replacement
 
     return [
         (
             _trace_and_lower_to_edge_ops(embedding_byte_pattern),
-            _trace_and_lower_to_edge_ops(embedding_byte_replacement),
-            [],
+            _trace_and_lower_to_edge_ops(get_embedding_byte_replacement(False)),
+            [get_embedding_replacement_filter(has_nonzero_zero_point=False)],
+        ),
+        (
+            _trace_and_lower_to_edge_ops(embedding_byte_pattern),
+            _trace_and_lower_to_edge_ops(get_embedding_byte_replacement(True)),
+            [get_embedding_replacement_filter(has_nonzero_zero_point=True)],
         ),
         (
             _trace_and_lower_to_edge_ops(embedding_byte_dtype_pattern),
-            _trace_and_lower_to_edge_ops(embedding_byte_dtype_replacement),
-            [],
+            _trace_and_lower_to_edge_ops(get_embedding_byte_dtype_replacement(False)),
+            [get_embedding_replacement_filter(has_nonzero_zero_point=False)],
+        ),
+        (
+            _trace_and_lower_to_edge_ops(embedding_byte_dtype_pattern),
+            _trace_and_lower_to_edge_ops(get_embedding_byte_dtype_replacement(True)),
+            [get_embedding_replacement_filter(has_nonzero_zero_point=True)],
         ),
         (
             _trace_and_lower_to_edge_ops(embedding_2bit_pattern),
-            _trace_and_lower_to_edge_ops(embedding_2bit_replacement),
-            [],
+            _trace_and_lower_to_edge_ops(get_embedding_2bit_replacement(False)),
+            [get_embedding_replacement_filter(has_nonzero_zero_point=False)],
+        ),
+        (
+            _trace_and_lower_to_edge_ops(embedding_2bit_pattern),
+            _trace_and_lower_to_edge_ops(get_embedding_2bit_replacement(True)),
+            [get_embedding_replacement_filter(has_nonzero_zero_point=True)],
         ),
         (
             _trace_and_lower_to_edge_ops(embedding_2bit_dtype_pattern),
-            _trace_and_lower_to_edge_ops(embedding_2bit_dtype_replacement),
-            [],
+            _trace_and_lower_to_edge_ops(get_embedding_2bit_dtype_replacement(False)),
+            [get_embedding_replacement_filter(has_nonzero_zero_point=False)],
+        ),
+        (
+            _trace_and_lower_to_edge_ops(embedding_2bit_dtype_pattern),
+            _trace_and_lower_to_edge_ops(get_embedding_2bit_dtype_replacement(True)),
+            [get_embedding_replacement_filter(has_nonzero_zero_point=True)],
         ),
         (
             _trace_and_lower_to_edge_ops(embedding_4bit_pattern),
-            _trace_and_lower_to_edge_ops(embedding_4bit_replacement),
-            [],
+            _trace_and_lower_to_edge_ops(get_embedding_4bit_replacement(has_nonzero_zero_point=False)),
+            [get_embedding_replacement_filter(has_nonzero_zero_point=False)],
+        ),
+        (
+            _trace_and_lower_to_edge_ops(embedding_4bit_pattern),
+            _trace_and_lower_to_edge_ops(get_embedding_4bit_replacement(has_nonzero_zero_point=True)),
+            [get_embedding_replacement_filter(has_nonzero_zero_point=True)],
         ),
         (
             _trace_and_lower_to_edge_ops(embedding_4bit_dtype_pattern),
-            _trace_and_lower_to_edge_ops(embedding_4bit_dtype_replacement),
-            [],
+            _trace_and_lower_to_edge_ops(get_embedding_4bit_dtype_replacement(has_nonzero_zero_point=False)),
+            [get_embedding_replacement_filter(has_nonzero_zero_point=False)],
+        ),
+         (
+            _trace_and_lower_to_edge_ops(embedding_4bit_dtype_pattern),
+            _trace_and_lower_to_edge_ops(get_embedding_4bit_dtype_replacement(has_nonzero_zero_point=True)),
+            [get_embedding_replacement_filter(has_nonzero_zero_point=True)],
         ),
     ]
 
@@ -1445,7 +1511,7 @@ def replacement(x, x_scale, x_zero_point, x_qmin, x_qmax):
 """
 
 
-def get_quant_patterns_and_replacements() -> (
+def get_quant_patterns_and_replacements(node_value_dict) -> (
     List[Tuple[Callable, Callable, List[Callable]]]
 ):
 
@@ -1457,6 +1523,6 @@ def get_quant_patterns_and_replacements() -> (
             *_get_slice_patterns_and_replacements(),
             # *_get_fixed_qparams_ops_patterns_and_replacements(),
             *_get_embedding_ops_patterns_and_replacements(),
-            *_get_embedding_ops_patterns_and_replacements_torchao(),
+            *_get_embedding_ops_patterns_and_replacements_torchao(node_value_dict),
         ]
     )