Assert quant_min/quant_max in embedding4bit (pytorch#7410)

* init

* format fix

Author: metascroy

examples/models/llama/source_transformation/quantize.py

@@ -729,18 +729,18 @@ def __init__(
     def forward(self, indices: torch.Tensor) -> torch.Tensor:
         if not self.packed:  # 8bit
             return torch.ops.quantized_decomposed.embedding_byte.dtype(
-                self.weight, self.scales, None, 0, 0, indices, dtype=self.dtype
+                self.weight, self.scales, None, -128, 127, indices, dtype=self.dtype
             )
         else:  # packed
             if self.bitwidth == 2:
                 return torch.ops.quantized_decomposed.embedding_2bit.dtype(
-                    self.weight, self.scales, None, 0, 0, indices, dtype=self.dtype
+                    self.weight, self.scales, None, -2, 1, indices, dtype=self.dtype
                 )
 
             # Remaining case (always return to make pyre happy)
             assert self.bitwidth == 4
             return torch.ops.quantized_decomposed.embedding_4bit.dtype(
-                self.weight, self.scales, None, 0, 0, indices, dtype=self.dtype
+                self.weight, self.scales, None, -8, 7, indices, dtype=self.dtype
             )
 
 

exir/passes/_quant_patterns_and_replacements.py

@@ -202,6 +202,13 @@ def embedding_2bit(
     weight_quant_max: int,
     indices: torch.Tensor,
 ) -> torch.Tensor:
+    assert (
+        weight_quant_min == -2
+    ), "embedding_2bit in ExecuTorch expects weight_quant_min == -2"
+    assert (
+        weight_quant_max == 1
+    ), "embedding_2bit in ExecuTorch expects weight_quant_max == 1"
+
     embedding_weight_checks(weight, weight_scales, weight_zero_points)
     group_size = (4 * weight.size(1)) // (
         weight_scales.size(1) if weight_scales.dim() == 2 else 1
@@ -257,6 +264,13 @@ def embedding_2bit_dtype(
     indices: torch.Tensor,
     dtype: Optional[torch.dtype],
 ) -> torch.Tensor:
+    assert (
+        weight_quant_min == -2
+    ), "embedding_2bit_dtype in ExecuTorch expects weight_quant_min == -2"
+    assert (
+        weight_quant_max == 1
+    ), "embedding_2bit_dtype in ExecuTorch expects weight_quant_max == 1"
+
     embedding_weight_checks(weight, weight_scales, weight_zero_points)
     group_size = (4 * weight.size(1)) // (
         weight_scales.size(1) if weight_scales.dim() == 2 else 1
@@ -334,6 +348,13 @@ def embedding_4bit(
     weight_quant_max: int,
     indices: torch.Tensor,
 ) -> torch.Tensor:
+    assert (
+        weight_quant_min == -8
+    ), "embedding_4bit in ExecuTorch expects weight_quant_min == -8"
+    assert (
+        weight_quant_max == 7
+    ), "embedding_4bit in ExecuTorch expects weight_quant_max == 7"
+
     embedding_weight_checks(weight, weight_scales, weight_zero_points)
     group_size = (2 * weight.size(1)) // (
         weight_scales.size(1) if weight_scales.dim() == 2 else 1
@@ -387,6 +408,13 @@ def embedding_4bit_dtype(
     indices: torch.Tensor,
     dtype: Optional[torch.dtype],
 ) -> torch.Tensor:
+    assert (
+        weight_quant_min == -8
+    ), "embedding_4bit_dtype in ExecuTorch expects weight_quant_min == -8"
+    assert (
+        weight_quant_max == 7
+    ), "embedding_4bit_dtype in ExecuTorch expects weight_quant_max == 7"
+
     embedding_weight_checks(weight, weight_scales, weight_zero_points)
     group_size = (2 * weight.size(1)) // (
         weight_scales.size(1) if weight_scales.dim() == 2 else 1