Move codebook (LUT) generation methods into common utils. Update functions be more compatible with coreml. (#2772)

Summary: Pull Request resolved: #2772

Reviewed By: metascroy

Differential Revision: D79595460

Author: szyszyzys

test/prototype/test_codebook_coreml.py

@@ -75,6 +75,41 @@ def test_quantize_api(self):
         )
         assert type(m[0].weight) == CodebookQuantizedTensor
 
+    def test_choose_qparams_codebook_row_grouping(self):
+        # Test with a block_size that forces row-wise grouping: [10, 256]
+        # Input tensor is (100, 256)
+        row_grouped_block_size = [10, -1]
+        num_row_groups = (
+            self.input.shape[0] // row_grouped_block_size[0]
+        )  # 100 // 10 = 10
+
+        codebook, wq = choose_qparams_and_quantize_codebook_coreml(
+            self.input,
+            self.code_dtype,
+            row_grouped_block_size,
+        )
+
+        # Expected shape for row-wise grouping is (num_row_groups, 1, 2**nbits, 1)
+        self.assertEqual(codebook.shape, (num_row_groups, 1, 2**self.nbits, 1))
+        self.assertEqual(wq.shape, (100, 256))
+
+        self.assertFalse(torch.isnan(codebook).any())
+        self.assertFalse(torch.isnan(wq).any())
+
+    def test_codebook_quantized_tensor_from_float_row_grouping(self):
+        # Test end-to-end quantization/dequantization with row grouping
+        row_grouped_block_size = [20, -1]  # 100 is divisible by 20
+        cqt = CodebookQuantizedTensor.from_float(
+            self.input,
+            self.code_dtype,
+            row_grouped_block_size,
+        )
+
+        dequant = cqt.dequantize()
+        # The SQNR will be different from column grouping, but should still be high
+        sqnr = compute_error(dequant, self.input)
+        self.assertGreater(sqnr, 30)
+
     def test_export(self):
         m = torch.nn.Sequential(torch.nn.Linear(128, 64)).to(torch.float32)
         quantize_(m, CodebookWeightOnlyConfig(self.code_dtype, self.block_size))

torchao/experimental/tests/test_groupwise_lowbit_weight_lut_quantizer.py

@@ -0,0 +1,170 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import copy
+import tempfile
+import unittest
+
+import torch
+import torch.nn as nn
+from parameterized import param, parameterized
+from torch import uint1, uint2, uint3, uint4
+
+from torchao.prototype.quantization.codebook_groupwise.api import (
+    GroupwiseLutWeightConfig,
+)
+from torchao.prototype.quantization.codebook_utils.codebook_utils import (
+    group_size_to_block_shapes,
+)
+from torchao.quantization.quant_api import quantize_
+
+
+class TestGroupwiseLowbitWeightLut(unittest.TestCase):
+    """
+    Test suite for the GroupwiseLutWeight quantization scheme, updated for the
+    new simplified API.
+    """
+
+    TEST_CASES = [
+        param(
+            code_dtype=code_dtype,
+            lut_group_size=lut_group_size,
+            weight_dtype=weight_dtype,
+            has_bias=has_bias,
+        )
+        for code_dtype in [uint1, uint2, uint3, uint4]
+        for lut_group_size in [256, 512]
+        for weight_dtype in [torch.float32]
+        for has_bias in [True, False]
+    ]
+
+    # --------------------------------------------------------------------------
+    # Test 1: End-to-End Model Accuracy
+    # --------------------------------------------------------------------------
+    @parameterized.expand(TEST_CASES)
+    def test_e2e_accuracy_vs_reference(
+        self,
+        code_dtype,
+        lut_group_size,
+        weight_dtype,
+        has_bias,
+    ):
+        """
+        Tests the numerical accuracy of the full quantized model against a reference.
+        This now uses the `use_qdq_reference` flag instead of layout objects.
+        """
+        m, k, n = 3, 64, 32
+        activations = torch.randn(m, k, dtype=weight_dtype)
+        model = nn.Sequential(nn.Linear(k, n, bias=has_bias)).to(dtype=weight_dtype)
+
+        # --- 2. Update tensor_shape to reflect the new (k, n) layout ---
+        lut_block_shape = group_size_to_block_shapes(
+            lut_group_size=lut_group_size, tensor_shape=(n, k)
+        )
+
+        # --- Quantize using C++ ops ---
+        quantized_model = copy.deepcopy(model)
+        perf_config = GroupwiseLutWeightConfig(
+            code_dtype=code_dtype,
+            weight_dtype=weight_dtype,
+            lut_block_shape=lut_block_shape,
+            use_qdq_reference=False,
+        )
+        quantize_(quantized_model, perf_config)
+        with torch.no_grad():
+            actual_result = quantized_model(activations)
+
+        # --- Quantize for Reference (using Python ops) ---
+        reference_model = copy.deepcopy(model)
+        ref_config = GroupwiseLutWeightConfig(
+            code_dtype=code_dtype,
+            weight_dtype=weight_dtype,
+            lut_block_shape=lut_block_shape,
+            use_qdq_reference=True,
+        )
+        quantize_(reference_model, ref_config)
+        with torch.no_grad():
+            expected_result = reference_model(activations)
+        # Compare results
+        self.assertTrue(
+            torch.allclose(actual_result, expected_result, atol=1e-2, rtol=1e-2)
+        )
+
+    def tearDown(self):
+        """
+        Clear the TorchDynamo cache after each test case to prevent
+        recompilation errors in parameterized tests.
+        """
+        super().tearDown()
+        torch._dynamo.reset()
+
+    # --------------------------------------------------------------------------
+    # Test 2: Deployment Readiness (Updated for new API)
+    # --------------------------------------------------------------------------
+    @parameterized.expand(TEST_CASES)
+    def test_export_compile_aoti(
+        self,
+        code_dtype,
+        lut_group_size,
+        weight_dtype,
+        has_bias,
+    ):
+        """
+        Tests that the quantized model can be exported and compiled.
+        """
+        k, n = 64, 32
+        activations = torch.randn(2, k, dtype=weight_dtype)
+        model = (
+            nn.Sequential(nn.Linear(k, n, bias=has_bias)).to(dtype=weight_dtype).eval()
+        )
+        lut_block_shape = group_size_to_block_shapes(
+            lut_group_size=lut_group_size,
+            tensor_shape=(n, k),
+        )
+
+        # Configure the quantization using the new API
+        config = GroupwiseLutWeightConfig(
+            code_dtype=code_dtype,
+            weight_dtype=weight_dtype,
+            lut_block_shape=lut_block_shape,
+            use_qdq_reference=False,
+        )
+        quantize_(model, config)
+
+        with torch.no_grad():
+            eager_results = model(activations)
+
+        # Export and Compile
+        exported_model = torch.export.export(model, (activations,))
+        compiled_model = torch.compile(model, fullgraph=True)
+
+        with tempfile.TemporaryDirectory() as tmpdir, torch.no_grad():
+            # Check exported model
+            exported_results = exported_model.module()(activations)
+            self.assertTrue(
+                torch.allclose(eager_results, exported_results, atol=1e-3, rtol=1e-3)
+            )
+
+            # Check compiled model
+            compiled_results = compiled_model(activations)
+            self.assertTrue(
+                torch.allclose(eager_results, compiled_results, atol=1e-3, rtol=1e-3)
+            )
+
+            # Check AOTI compiled model using the packaging API
+            package_path = f"{tmpdir}/model.pt2"
+            torch._inductor.aoti_compile_and_package(
+                exported_model, package_path=package_path
+            )
+            aoti_model = torch._inductor.aoti_load_package(package_path)
+            aoti_results = aoti_model(activations)
+            self.assertTrue(
+                torch.allclose(eager_results, aoti_results, atol=1e-3, rtol=1e-3)
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

torchao/prototype/quantization/codebook_coreml/codebook_ops.py

@@ -57,65 +57,83 @@ def choose_qparams_and_quantize_codebook_coreml(
     assert code_dtype in list(_SUB_BYTE_UINT_BOUNDS.keys()) + [torch.uint8]
     nbits = _DTYPE_TO_BIT_WIDTH[code_dtype]
     assert nbits >= 1 and nbits <= 8, f"nbits must be in [1, 8], got {nbits}"
-
-    assert len(block_size) == input_tensor.dim()
-    block_size = block_size.copy()
-    for i in range(len(block_size)):
-        if block_size[i] == -1:
-            block_size[i] = input_tensor.shape[i]
-        assert block_size[i] >= 1 and input_tensor.shape[i] % block_size[i] == 0, (
-            "block_size[i] must divide input_tensor.shape[i]"
-        )
-
     assert input_tensor.dim() == 2, "Currently only rank 2 tensors are supported"
-    assert block_size[0] == input_tensor.shape[0], (
-        "Currently only support per-grouped channel granularity"
-    )
     assert cluster_dim == 1, (
         f"only cluster_dim == 1 is supported right now, got {cluster_dim}"
     )
 
-    num_lut = input_tensor.shape[1] // block_size[1]
-    group_size = block_size[1]
-
-    # for converting to numpy
-    input_tensor = input_tensor.detach()
     original_shape = input_tensor.shape
+    N, K = original_shape
+    input_tensor = input_tensor.detach()
 
-    # reshape to (N, K // group_size, group_size)
-    input_tensor = input_tensor.reshape(input_tensor.shape[0], num_lut, group_size)
-    from coremltools.models.neural_network.quantization_utils import (
-        _get_kmeans_lookup_table_and_weight,
+    # --- Process block_size ---
+    assert len(block_size) == 2
+    processed_block_size = block_size.copy()
+    if processed_block_size[0] == -1:
+        processed_block_size[0] = N
+    if processed_block_size[1] == -1:
+        processed_block_size[1] = K
+
+    row_block_size, col_block_size = processed_block_size
+    assert N % row_block_size == 0, (
+        f"Tensor rows ({N}) not divisible by row block size ({row_block_size})"
+    )
+    assert K % col_block_size == 0, (
+        f"Tensor cols ({K}) not divisible by col block size ({col_block_size})"
     )
 
-    res_lut = []
-    # each res_w[:, i, :] will use the same lookup table
-    # res_w: (N, K // group_size, group_size)
-    res_w = torch.zeros_like(input_tensor, dtype=torch.uint8)
-    for i in range(num_lut):
-        # lut: (2**nbits, 1)
-        # w: (N * group_size)
-        lut, w = _get_kmeans_lookup_table_and_weight(
-            nbits, input_tensor[:, i, :], force_kmeans1d, cluster_dim, vector_axis
-        )
-        res_lut.append(torch.from_numpy(lut))
-        res_w[:, i, :] = torch.from_numpy(w.reshape(input_tensor.shape[0], group_size))
-
-    # directly stack all lookup tables along dim 0
-    # res_lut: (K // group_size, 2 ** nbits)
-    res_lut = torch.stack(res_lut, dim=0)
-
-    # The final LUT should have dimension equal to input_tensor.dim() + 2
-    # The first input_tensor.dim() dimensions index over the tables,
-    # input_tensor.dim() + 1 indexes over the nbit indices
-    # input_tensor.dim() + 2 are the look up values (shape = 1 for scalar)
-    # res_lut: (N, K // group_size, 2 ** nbits, group_size)
-    res_lut = res_lut.reshape(1, num_lut, 2**nbits, 1)
+    # --- Determine and execute grouping strategy ---
+    assert row_block_size == N or col_block_size == K
+    is_col_grouping = row_block_size == N
 
-    # reshape back to (N, K)
-    res_w = res_w.reshape(*original_shape)
+    res_lut_list = []
+    from coremltools.models.neural_network.quantization_utils import (
+        _get_kmeans_lookup_table_and_weight,
+    )
 
-    return res_lut, res_w
+    if is_col_grouping:
+        # STRATEGY 1: Group by COLUMNS
+        num_luts = K // col_block_size
+        reshaped_tensor = input_tensor.reshape(N, num_luts, col_block_size)
+        res_codes = torch.zeros_like(reshaped_tensor, dtype=torch.uint8)
+
+        for i in range(num_luts):
+            block_to_quantize = reshaped_tensor[:, i, :]
+            lut, w = _get_kmeans_lookup_table_and_weight(
+                nbits, block_to_quantize, force_kmeans1d, cluster_dim, vector_axis
+            )
+            res_lut_list.append(torch.from_numpy(lut))
+            res_codes[:, i, :] = torch.from_numpy(w.reshape(N, col_block_size))
+
+        # Shape to match CoreML spec: (1, num_luts, 2**nbits, 1)
+        final_luts = torch.stack(res_lut_list, dim=0).reshape(1, num_luts, 2**nbits, 1)
+
+    else:  # is_row_grouping
+        # STRATEGY 2: Group by ROWS
+        num_luts = N // row_block_size
+        reshaped_tensor = input_tensor.reshape(num_luts, row_block_size, K)
+        res_codes = torch.zeros_like(reshaped_tensor, dtype=torch.uint8)
+
+        for i in range(num_luts):
+            block_to_quantize = reshaped_tensor[i, :, :]
+            lut, w = _get_kmeans_lookup_table_and_weight(
+                nbits, block_to_quantize, force_kmeans1d, cluster_dim, vector_axis
+            )
+            res_lut_list.append(torch.from_numpy(lut))
+            res_codes[i, :, :] = torch.from_numpy(w.reshape(row_block_size, K))
+
+        final_luts_stacked = torch.stack(
+            res_lut_list, dim=0
+        )  # Shape: (num_luts, 2**nbits, 1)
+
+        # Reshape to the consistent 4D format
+        # The shape is (num_row_groups, 1, 2**nbits, 1)
+        final_luts = final_luts_stacked.reshape(num_luts, 1, 2**nbits, 1)
+
+    # Reshape codes back to the original tensor shape
+    final_codes = res_codes.reshape(*original_shape)
+
+    return final_luts, final_codes
 
 
 @register_custom_op

torchao/prototype/quantization/codebook_groupwise/__init__.py

@@ -1,4 +1,4 @@
 from .api import GroupwiseLutWeightConfig
-from .codebook_quantized_tensor import GroupwiseLutQuantizedTensor
+from .codebook_quantized_tensor import CodebookQuantizedPackedTensor
 
-__all__ = ["GroupwiseLutQuantizedTensor", "GroupwiseLutWeightConfig"]
+__all__ = ["CodebookQuantizedPackedTensor", "GroupwiseLutWeightConfig"]