Update quant_primitives.py

Author: iseeyuan

test/quantization/test_quant_primitives.py

@@ -164,7 +164,7 @@ def test_get_group_qparams_symmetric(self):
         scale_obs = scale_obs.reshape(weight.shape[0], -1)
 
         # assert that scales are identical
-        (scale_ao, _) = get_group_qparams_symmetric(weight, n_bit, groupsize, precision=torch.float16)
+        (scale_ao, _) = get_group_qparams_symmetric(weight, n_bit, groupsize, precision=torch.float16, mapping_type=MappingType.SYMMETRIC)
         torch.testing.assert_close(scale_obs, scale_ao, rtol=0, atol=0)
 
     def test_choose_qparams_group_sym(self):
@@ -179,11 +179,27 @@ def test_choose_qparams_group_sym(self):
         precision = torch.float32
         scale, zero_point = choose_qparams_affine(input, mapping_type, block_size, dtype, eps=eps, scale_dtype=precision, zero_point_dtype=precision)
 
-        scale_ref, zp_ref = get_group_qparams_symmetric(input, n_bit=8, groupsize=2, precision=precision)
+        scale_ref, zp_ref = get_group_qparams_symmetric(input, n_bit=8, groupsize=2, precision=precision, mapping_type=mapping_type)
 
         self.assertTrue(torch.equal(scale, scale_ref))
         self.assertTrue(torch.equal(zero_point, zp_ref))
 
+    def test_choose_qparams_group_sym_pos_neg(self):
+        """
+        Test the added MappingType.SYMMETRIC_MAX_POS_NEG
+        """
+        input = torch.randn(10, 10)
+        mapping_type = MappingType.SYMMETRIC_MAX_POS_NEG
+        dtype = torch.int8
+        block_size = (1, 2)
+        eps = torch.finfo(torch.float32).eps
+        precision = torch.float32
+        scale, zero_point = choose_qparams_affine(input, mapping_type, block_size, dtype, eps=eps, scale_dtype=precision, zero_point_dtype=precision)
+
+        scale_ref, zp_ref = get_group_qparams_symmetric(input, n_bit=8, groupsize=2, precision=precision, mapping_type=mapping_type)
+
+        self.assertTrue(torch.equal(scale, scale_ref))
+        self.assertTrue(torch.equal(zero_point, zp_ref))
     @unittest.skipIf(not TORCH_VERSION_AT_LEAST_2_3, "skipping when torch version is 2.3 or lower")
     @unittest.skipIf(is_fbcode(), "broken in fbcode")
     def test_choose_qparams_token_asym(self):

torchao/quantization/quant_primitives.py

@@ -41,12 +41,18 @@ class MappingType(Enum):
     we'll use (-10.2, 10.2) as the range for floating point and map that to (-8, 7)
     e.g. scale = (10.2 - (-10.2)) / (7 - (-8))
 
+    SYMMETRIC_MAX_POS_NEG is a variant of symmetric mapping, where the scale is the max of smin
+    and smax, where smin = min_val_neg / quant_min, and smax = max_val_pos / quant_max. By calculating
+    smin and smax individually, there can be less round error on negative values, and no out-of-range
+    of all floating point values.
+
     asymmetric mapping means we just directly map the floating point range to integer range,
     for the above example, we will map (-3.5, 10.2) to (-8, 7) and calculate quantization parameter
     based on this mapping
     e.g. scale = (10.2 - (-3.5)) / (7 - (-8))
     """
     SYMMETRIC = auto()
+    SYMMETRIC_MAX_POS_NEG = auto()
     ASYMMETRIC = auto()
 
 class ZeroPointDomain(Enum):
@@ -695,7 +701,7 @@ def _choose_qparams_affine(
        and `zero_point_domain`
     """
     quant_min, quant_max = _get_and_check_qmin_qmax(target_dtype, quant_min, quant_max)
-    assert mapping_type in [MappingType.SYMMETRIC.name, MappingType.ASYMMETRIC.name], f"Unsupported mapping type: {mapping_type}"
+    assert mapping_type in [MappingType.SYMMETRIC.name, MappingType.SYMMETRIC_MAX_POS_NEG.name, MappingType.ASYMMETRIC.name], f"Unsupported mapping type: {mapping_type}"
 
     if input is not None:
         if scale_dtype is None:
@@ -729,11 +735,25 @@ def _choose_qparams_affine(
         min_val_neg = min_val
         max_val_pos = max_val
 
-    if mapping_type == MappingType.SYMMETRIC.name:
-        smin = min_val_neg / float(quant_min)
-        smax = max_val_pos / float(quant_max)
-        mask = smin > smax
-        scale = torch.where(mask, smin, smax)
+    if mapping_type == MappingType.SYMMETRIC.name or mapping_type == MappingType.SYMMETRIC_MAX_POS_NEG.name:
+        # scales
+        if mapping_type == MappingType.SYMMETRIC.name:
+            max_val_pos = torch.max(-min_val_neg, max_val_pos)
+            scale = max_val_pos / (float(quant_max - quant_min) / 2)
+        else:
+            assert mapping_type == MappingType.SYMMETRIC_MAX_POS_NEG.name
+            # calculate smin and smax individually and choose the larger one. For example, if quant_min = -8 and
+            # quant_max = 7.
+            # - If smin is bigger: There would be coverage on negative values down to -8, and less rounding
+            # error than the existing SYMMETRIC case.
+            # - If smax is bigger: it covers the positive values up to 7. The round
+            # error may be bigger than the existing SYMMETRIC case. Either way, there's no out-of-range fp values after
+            # quantization.
+            smin = min_val_neg / float(quant_min)
+            smax = max_val_pos / float(quant_max)
+            mask = smin > smax
+            scale = torch.where(mask, smin, smax)
+        # zeros
         if not preserve_zero:
             raise ValueError("preserve_zero == False is not supported for symmetric quantization")
         if zero_point_domain is not None and zero_point_domain != ZeroPointDomain.INT.name:

torchao/quantization/utils.py

@@ -418,7 +418,7 @@ def groupwise_affine_dequantize_tensor(
 
 
 # TODO: separate scale and zero point precision
-def get_group_qparams_symmetric(w, n_bit=4, groupsize=128, precision=torch.float32):
+def get_group_qparams_symmetric(w, n_bit=4, groupsize=128, precision=torch.float32, mapping_type=MappingType.SYMMETRIC_MAX_POS_NEG):
     # needed for GPTQ with padding
     if groupsize > w.shape[-1]:
         groupsize = w.shape[-1]
@@ -427,7 +427,6 @@ def get_group_qparams_symmetric(w, n_bit=4, groupsize=128, precision=torch.float
     assert w.dim() == 2
     assert n_bit <= 8, f"unsupported n_bit: {n_bit}"
 
-    mapping_type = MappingType.SYMMETRIC
     block_size = (1, groupsize)
     eps = torch.finfo(torch.float32).eps
     ranges = {}
@@ -445,8 +444,9 @@ def group_quantize_tensor_symmetric(
     n_bit=4,
     group_size=128,
     precision=torch.float32,
+    mapping_type=MappingType.SYMMETRIC_MAX_POS_NEG
 ):
-    scales, zeros = get_group_qparams_symmetric(w, n_bit, group_size, precision)
+    scales, zeros = get_group_qparams_symmetric(w, n_bit, group_size, precision, mapping_type)
     n_bit = 4
     max_int = 2 ** (n_bit - 1) - 1
     min_int = -(2 ** (n_bit - 1))