Add IntxUnpackedTensor (#2732)

* add intx unpacked tensor

* up

* up

* up

* up

* up

Author: metascroy

test/quantization/quantize_/workflows/intx/test_intx_unpacked_tensor.py

@@ -0,0 +1,147 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+import torch
+from torch.testing._internal.common_utils import (
+    TestCase,
+    run_tests,
+)
+
+from torchao.quantization import (
+    IntxWeightOnlyConfig,
+    quantize_,
+)
+from torchao.quantization.granularity import PerGroup
+from torchao.quantization.utils import compute_error
+from torchao.utils import (
+    TORCH_VERSION_AT_LEAST_2_8,
+)
+
+
+@unittest.skipIf(not TORCH_VERSION_AT_LEAST_2_8, "Need pytorch 2.8+")
+class TestIntxUnpackedTensor(TestCase):
+    def setUp(self):
+        self.config = IntxWeightOnlyConfig(
+            weight_dtype=torch.int4,
+            granularity=PerGroup(32),
+            version=2,
+        )
+
+    def test_embedding(self):
+        dtype = torch.bfloat16
+        device = "cpu"
+        input = torch.randint(low=0, high=128, size=(10,), device=device)
+        embedding = torch.nn.Embedding(128, 256, dtype=dtype, device=device)
+        original = embedding(input)
+        quantize_(embedding, self.config)
+        quantized = embedding(input)
+        error = compute_error(original, quantized)
+        self.assertTrue(error > 20)
+
+    def test_linear(self):
+        dtype = torch.bfloat16
+        device = "cpu"
+        input = torch.randn(1, 128, dtype=dtype, device=device)
+        linear = torch.nn.Linear(128, 256, dtype=dtype, device=device)
+        original = linear(input)
+        quantize_(linear, self.config)
+        quantized = linear(input)
+        error = compute_error(original, quantized)
+        self.assertTrue(error > 20)
+
+    def test_slice(self):
+        dtype = torch.bfloat16
+        device = "cpu"
+        dummy = torch.nn.Linear(256, 256, bias=False, dtype=dtype, device=device)
+
+        dummy1 = torch.nn.Linear(256, 64, bias=False, dtype=dtype, device=device)
+        dummy1.weight = torch.nn.Parameter(
+            dummy.weight.narrow(0, 0, 64), requires_grad=False
+        )
+
+        dummy2 = torch.nn.Linear(128, 256, dtype=dtype, device=device)
+        dummy2.weight = torch.nn.Parameter(
+            dummy.weight.narrow(1, 0, 128), requires_grad=False
+        )
+
+        quantize_(dummy, self.config)
+        weight1 = dummy.weight.narrow(0, 0, 64)
+        weight2 = dummy.weight.narrow(1, 0, 128)
+
+        self.assertEqual(weight1.qdata, dummy.weight.qdata.narrow(0, 0, 64))
+        self.assertEqual(weight1.scale, dummy.weight.scale.narrow(0, 0, 64))
+
+        self.assertEqual(weight2.qdata, dummy.weight.qdata.narrow(1, 0, 128))
+        self.assertEqual(weight2.scale, dummy.weight.scale.narrow(1, 0, 4))
+
+        # check for sliced weight, before and after float8 quantization
+        # does not differ too much
+        input = torch.randn(2, 256, dtype=dtype, device=device)
+        res_ref = dummy1(input)
+        dummy.weight = torch.nn.Parameter(weight1, requires_grad=False)
+        res = dummy(input)
+        assert compute_error(res, res_ref) > 20
+
+        input = torch.randn(2, 128, dtype=dtype, device=device)
+        res_ref = dummy2(input)
+        dummy.weight = torch.nn.Parameter(weight2, requires_grad=False)
+        res = dummy(input)
+        assert compute_error(res, res_ref) > 15
+
+    def test_slice_and_copy_(self):
+        device = "cpu"
+        l = torch.nn.Linear(1024, 1024).to(device).to(torch.bfloat16)
+        l.weight = torch.nn.Parameter(
+            torch.zeros(1024, 1024, dtype=torch.bfloat16, device=device)
+        )
+        quantize_(l, self.config)
+        param = l.weight
+        param_data = param.data
+        param_data = param_data.narrow(0, 0, 512)
+        assert param.data.qdata.data_ptr() == param_data.qdata.data_ptr()
+        assert param.data.scale.data_ptr() == param_data.scale.data_ptr()
+        assert param.data.zero_point.data_ptr() == param_data.zero_point.data_ptr()
+        orig_value = param.data.qdata[0][0].item()
+
+        # dummy_l has random input (shouldn't be 0)
+        dummy_l = torch.nn.Linear(1024, 1024).to(device).to(torch.bfloat16)
+        quantize_(dummy_l, self.config)
+        quantized = dummy_l.weight
+        quantized = quantized.narrow(0, 0, 512)
+
+        param_data.copy_(quantized)
+
+        # making sure param.data is updated
+        assert param.data.qdata[0][0] != orig_value
+
+    def test_to_dtype(self):
+        activations_bf16 = torch.randn(1, 128, dtype=torch.bfloat16)
+        activations_fp32 = torch.randn(1, 128, dtype=torch.float32)
+        activations_fp16 = torch.randn(1, 128, dtype=torch.float16)
+
+        linear = torch.nn.Linear(128, 256)
+        quantize_(linear, self.config)
+
+        linear.to(dtype=torch.float16)
+        linear(activations_fp16)
+
+        linear.to(dtype=torch.float32)
+        linear(activations_fp32)
+
+        linear.to(dtype=torch.bfloat16)
+        linear(activations_bf16)
+
+    def test_export(self):
+        linear = torch.nn.Linear(128, 256)
+        quantize_(linear, self.config)
+        ep = torch.export.export(linear, (torch.randn(1, 128),))
+        assert "torch.ops.torchao.dequantize_affine.default" in ep.graph_module.code
+
+
+if __name__ == "__main__":
+    run_tests()

torchao/quantization/__init__.py

@@ -93,6 +93,7 @@
     Int4MarlinSparseTensor,
     Int4PreshuffledTensor,
     Int4Tensor,
+    IntxUnpackedTensor,
 )
 from .smoothquant import (
     SmoothFakeDynamicallyQuantizedLinear,
@@ -161,6 +162,7 @@
     "Int4Tensor",
     "Int4PreshuffledTensor",
     "Int4MarlinSparseTensor",
+    "IntxUnpackedTensor",
     "Float8Tensor",
     # smooth quant - subject to change
     "get_scale",

torchao/quantization/quant_api.py

@@ -75,6 +75,7 @@
     Int4MarlinSparseTensor,
     Int4PreshuffledTensor,
     Int4Tensor,
+    IntxUnpackedTensor,
     QuantizeTensorToFloat8Kwargs,
 )
 from torchao.quantization.transform_module import (
@@ -454,6 +455,10 @@ def _linear_extra_repr(self):
     return f"in_features={self.weight.shape[1]}, out_features={self.weight.shape[0]}, weight={_quantization_type(self.weight)}"
 
 
+def _embedding_extra_repr(self):
+    return f"num_embeddings={self.weight.shape[0]}, embedding_dim={self.weight.shape[1]}, weight={_quantization_type(self.weight)}"
+
+
 def _get_linear_subclass_inserter(
     constructor, *, allow_requires_grad=False, propagate_bias=False, **kwargs
 ):
@@ -1987,6 +1992,8 @@ class IntxWeightOnlyConfig(AOBaseConfig):
     mapping_type: MappingType = MappingType.SYMMETRIC
     scale_dtype: Optional[torch.dtype] = None
     layout: Layout = QDQLayout()
+    packing_format: PackingFormat = PackingFormat.UNPACKED_TO_INT8
+    version: int = 1
 
     def __post_init__(self):
         torch._C._log_api_usage_once("torchao.quantization.IntxWeightOnlyConfig")
@@ -2005,16 +2012,13 @@ def __post_init__(self):
         )
 
 
-@register_quantize_module_handler(IntxWeightOnlyConfig)
-def _intx_weight_only_transform(
-    module: torch.nn.Module, config: IntxWeightOnlyConfig
-) -> torch.nn.Module:
-    weight = module.weight
+def _intx_weight_only_quantize_tensor(weight, config):
     weight_dtype = config.weight_dtype
     granularity = config.granularity
     mapping_type = config.mapping_type
     scale_dtype = config.scale_dtype
     layout = config.layout
+    packing_format = config.packing_format
 
     assert weight.dim() == 2, (
         f"IntxWeightOnlyConfig only works for 2-d Tensor, got: {weight.dim()}"
@@ -2029,11 +2033,28 @@ def _intx_weight_only_transform(
     else:
         raise ValueError(f"granularity must be PerGroup or PerAxis, got {granularity}")
 
+    block_size = (1, group_size)
+
+    if config.version == 2:
+        if config.packing_format == PackingFormat.UNPACKED_TO_INT8:
+            new_weight = IntxUnpackedTensor.from_hp(
+                weight,
+                block_size,
+                weight_dtype,
+                mapping_type=mapping_type,
+            )
+            if scale_dtype is not None and scale_dtype != weight.dtype:
+                new_weight.scale = new_weight.scale.to(scale_dtype).to(weight.dtype)
+            return new_weight
+        else:
+            raise ValueError(f"Unsupported packing format: {packing_format}")
+
+    # Version 1
     quant_min, quant_max = _DTYPE_TO_QVALUE_BOUNDS[weight_dtype]
     weight = to_affine_quantized_intx(
         input_float=weight,
         mapping_type=mapping_type,
-        block_size=(1, group_size),
+        block_size=block_size,
         target_dtype=torch.int8,
         quant_min=quant_min,
         quant_max=quant_max,
@@ -2043,7 +2064,25 @@ def _intx_weight_only_transform(
         zero_point_domain=ZeroPointDomain.INT,
         _layout=layout,
     )
-    module.weight = torch.nn.Parameter(weight, requires_grad=False)
+    return weight
+
+
+@register_quantize_module_handler(IntxWeightOnlyConfig)
+def _intx_weight_only_transform(
+    module: torch.nn.Module, config: IntxWeightOnlyConfig
+) -> torch.nn.Module:
+    assert hasattr(module, "weight"), (
+        "applying intx weight only quant requires module to have weight attribute"
+        + " but {module} does not have one"
+    )
+    new_weight = _intx_weight_only_quantize_tensor(module.weight, config)
+    module.weight = torch.nn.Parameter(new_weight, requires_grad=False)
+
+    if isinstance(module, nn.Linear):
+        module.extra_repr = types.MethodType(_linear_extra_repr, module)
+    elif isinstance(module, nn.Embedding):
+        module.extra_repr = types.MethodType(_embedding_extra_repr, module)
+
     return module
 
 

torchao/quantization/quantize_/common/packing_format.py

@@ -35,3 +35,8 @@ class PackingFormat(str, Enum):
     marlin_sparse is referring to the format used by marlin kernels, only supports symmetric quantization
     """
     MARLIN_SPARSE = "marlin_sparse"
+
+    """
+    Unpacked means the subbyte quantized data is stored as int8
+    """
+    UNPACKED_TO_INT8 = "unpacked_to_int8"

torchao/quantization/quantize_/workflows/__init__.py

@@ -11,11 +11,15 @@
 from .int4.int4_tensor import (
     Int4Tensor,
 )
+from .intx.intx_unpacked_tensor import (
+    IntxUnpackedTensor,
+)
 
 __all__ = [
     "Int4Tensor",
     "Int4PreshuffledTensor",
     "Int4MarlinSparseTensor",
     "Float8Tensor",
     "QuantizeTensorToFloat8Kwargs",
+    "IntxUnpackedTensor",
 ]

torchao/quantization/quantize_/workflows/intx/__init__.py

@@ -0,0 +1,5 @@
+from .intx_unpacked_tensor import IntxUnpackedTensor
+
+__all__ = [
+    "IntxUnpackedTensor",
+]