Add codebook support to CoreML using quantize_

Author: metascroy

backends/apple/coreml/compiler/torch_ops.py

@@ -8,149 +8,7 @@
 # coremltools than is used by ExecuTorch.  Each op registered here should have a link to a PR in coremltools that adds
 # the op to the coremltools library.
 
-import torch as _torch
-from coremltools import _logger
-from coremltools.converters.mil.frontend import _utils
-from coremltools.converters.mil.frontend.torch.ops import (
-    _get_inputs,
-    _get_kwinputs,
-    NUM_TO_NUMPY_DTYPE,
-    NUM_TO_TORCH_DTYPE,
-    split,
-    to,
-    transpose,
-    unbind,
-)
 import numpy as np
-from coremltools.converters.mil.frontend.torch.torch_op_registry import (
-    register_torch_op,
-)
-from coremltools.converters.mil.mil import types
-from executorch.exir.dim_order_utils import get_memory_format
-
-
-# https://github.com/apple/coremltools/pull/2556
-@register_torch_op(override=False)
-def transpose_copy(context, node):
-    transpose(context, node)
-
-
-# https://github.com/apple/coremltools/pull/2557
-@register_torch_op(override=False)
-def unbind_copy(context, node):
-    unbind(context, node)
-
-
-# https://github.com/apple/coremltools/pull/2563
-@register_torch_op(override=False)
-def split_copy(context, node):
-    split(context, node)
-
-
-@register_torch_op(
-    torch_alias=[
-        "dim_order_ops::_to_dim_order_copy",
-        "dim_order_ops._to_dim_order_copy",
-    ],
-    override=False,
-)
-def _to_dim_order_copy(context, node):
-    dim_order = _get_kwinputs(context, node, "dim_order", default=[None])[0]
-    node.kwinputs.pop("dim_order")
-
-    # In CoreML, dim_order.val will be an ndarray, so we convert it to a list
-    dim_order = [int(d) for d in dim_order.val]
-    memory_format = get_memory_format(dim_order)
-    assert (
-        memory_format == _torch.contiguous_format
-    ), "Only contiguous memory format is supported in CoreML"
-    to(context, node)
-
-
-# https://github.com/apple/coremltools/pull/2558
-@register_torch_op(
-    torch_alias=["torchao::dequantize_affine", "torchao.dequantize_affine"],
-    override=False,
-)
-def dequantize_affine(context, node):
-    inputs = _get_inputs(context, node, expected=[7, 8])
-    int_data = inputs[0].val
-    block_size = inputs[1].val
-    scale = inputs[2].val
-    zero_point = (
-        inputs[3].val if inputs[3] is not None and inputs[3].val is not None else None
-    )
-    # I do not think we need to worry about input_dtype b/c it gets cast to int4/int8
-    # For now, we just check that it is int8 or int32
-    input_dtype = inputs[4].val  # noqa: F841
-    assert NUM_TO_TORCH_DTYPE[input_dtype] in [
-        _torch.int8,
-        _torch.int32,
-    ], "input_dtype should be int8 or int32"
-
-    quant_min = inputs[5].val
-    quant_max = inputs[6].val
-
-    assert len(int_data.shape) == 2, "dequantize_affine only supports rank 2 inputs"
-
-    assert len(int_data.shape) == len(
-        block_size
-    ), "block_size must have the same length as int_data.shape"
-    assert block_size[0] == 1, "block_size[0] must be 1"
-    group_size = block_size[1]
-    k = int_data.shape[1]
-    assert k % group_size == 0, "k must be divisible by group_size"
-    scales_per_row = k // group_size
-    scale = scale.reshape(-1, scales_per_row)
-    if zero_point is not None:
-        zero_point = zero_point.reshape(-1, scales_per_row)
-
-    # TODO: I don't know if CoreML can make use of this
-    # We could add a cast op to the output, but I'm pretty CoreML will remove this during a later pass
-    # For now, we just log a warning
-    out_np_dtype = None
-    if len(inputs) > 7:
-        out_np_dtype = NUM_TO_NUMPY_DTYPE[inputs[7].val]
-        _logger.warning(
-            f"Core ML ignores output_dtype {out_np_dtype} on torchao.dequantize_affine and instead uses the native precision."
-        )
-
-    if quant_min == -8 and quant_max == 7:
-        quantized_np_dtype = types.nptype_from_builtin(types.string_to_builtin("int4"))
-    elif quant_min == -128 and quant_max == 127:
-        quantized_np_dtype = types.nptype_from_builtin(types.string_to_builtin("int8"))
-    else:
-        raise ValueError(
-            f"Unsupported quantization range: {quant_min} to {quant_max}.  CoreML only supports 4-bit and 8-bit quantization."
-        )
-
-    output = _utils._construct_constexpr_dequant_op(
-        int_data.astype(quantized_np_dtype),
-        zero_point,
-        scale,
-        axis=-1,
-        name=node.name,
-    )
-    context.add(output, node.name)
-
-
-
-# codes: torch.Tensor,
-#     codebook: torch.Tensor,
-#     code_dtype: torch.dtype,
-#     block_size: List[int],
-#     output_dtype: torch.dtype = torch.float32,
-
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
-
-# This file registers torch ops that are not yet in coremltools, or are in a more recent version of
-# coremltools than is used by ExecuTorch.  Each op registered here should have a link to a PR in coremltools that adds
-# the op to the coremltools library.
-
 import torch as _torch
 from coremltools import _logger
 from coremltools.converters.mil.frontend import _utils
@@ -164,8 +22,6 @@ def dequantize_affine(context, node):
     transpose,
     unbind,
 )
-import numpy as np
-
 from coremltools.converters.mil.frontend.torch.torch_op_registry import (
     register_torch_op,
 )
@@ -283,24 +139,25 @@ def dequantize_affine(context, node):
     override=False,
 )
 def dequantize_codebook(context, node):
-    print("IN DEQUANTIZE CODEBOOK")
     inputs = _get_inputs(context, node, expected=[4, 5])
     codes = inputs[0].val
     codebook = inputs[1].val
-    code_dtype = inputs[2].val
-    block_size = inputs[3].val
+    nbits = inputs[2].val
 
-    
-    assert len(codes.shape) == 2, "Only rank 2 inputs are supported"
+    # information in block_size is redundant with codebook.shape
+    block_size = inputs[3].val  # noqa: F841
 
-    # In TorchAO, the codebook shape is (n_lut, nbit, 1).  The LUTs are for the columns.
-    # In CoreML, the expected shape is (lut_block_size, nbit, 1).  1 here is for scalar
-    # lut_block_size has the same rank as codes/idxs and tells you how many LUTs there are per block, e.g., 
-    # lut_block_size=(1, 8) means there is 1 LUT per 8 columns
-    assert len(codebook.shape) == 3, "Only rank 3 inputs are supported for codebook"
-    assert codebook.shape[-1] == 1, "we only support scalar palletization"
-    codebook = np.expand_dims(codebook, 0)
+    assert len(codes.shape) == 2, "Only rank 2 inputs are supported"
 
+    # Assert codebook is as expected.  codebook.dim() = codes.dim() + 2
+    assert len(codebook.shape) == 4, "Only rank 4 inputs are supported for codebook"
+    assert codebook.shape[0] == 1, "Only grouped_channel granularity is supported"
+    n_luts = codebook.shape[1]
+    assert (
+        codes.shape[1] % n_luts == 0
+    ), "codes.shape[1] must be divisible by codebook.shape[1]"
+    assert codebook.shape[2] == 2**nbits
+    assert codebook.shape[3] == 1, "Only scalar look up values are supported"
 
     if len(inputs) > 4:
         output_dtype = inputs[4].val
@@ -309,11 +166,6 @@ def dequantize_codebook(context, node):
             f"Core ML ignores output_dtype {out_np_dtype} on torchao.dequantize_affine and instead uses the native precision."
         )
 
-    print("codes", codes.shape)
-    print("codebook", codebook.shape)
-    print("code_dtype", code_dtype)
-    print("block_size", block_size)
-
     output = _utils._construct_constexpr_lut_op(
         codes.astype(np.int8),
         codebook,

backends/apple/coreml/test/test_torch_ops.py

@@ -14,11 +14,9 @@
 
 from executorch.backends.apple.coreml.compiler import CoreMLBackend
 from executorch.backends.apple.coreml.partition import CoreMLPartitioner
-from torchao.quantization import IntxWeightOnlyConfig, PerAxis, PerGroup, quantize_
 
-from torchao.prototype.quantization.codebook_coreml import (
-    CodebookWeightOnlyConfig,
-)
+from torchao.prototype.quantization.codebook_coreml import CodebookWeightOnlyConfig
+from torchao.quantization import IntxWeightOnlyConfig, PerAxis, PerGroup, quantize_
 
 
 def is_fbcode():
@@ -167,12 +165,12 @@ def test_dequantize_affine_c8w_embedding_b4w_linear(self):
                 ], f"Got unexpected node target after delegation: {node.target.__name__}"
         et_prog = delegated_program.to_executorch()
         self._compare_outputs(et_prog, model, example_inputs)
-    
+
     def test_dequantize_codebook_linear(self):
         model, example_inputs = self._get_test_model()
         quantize_(
             model,
-            CodebookWeightOnlyConfig(dtype=torch.uint8, block_size=[-1, 16]),
+            CodebookWeightOnlyConfig(dtype=torch.uint2, block_size=[-1, 16]),
         )
         ep = torch.export.export(model, example_inputs)
         print("ORIGINAL MODEL", ep)
@@ -190,12 +188,35 @@ def test_dequantize_codebook_linear(self):
         print(et_prog.exported_program())
         self._compare_outputs(et_prog, model, example_inputs)
 
+    def test_dequantize_codebook_embedding(self):
+        model, example_inputs = self._get_test_model()
+        quantize_(
+            model,
+            CodebookWeightOnlyConfig(dtype=torch.uint3, block_size=[-1, 16]),
+            lambda m, fqn: isinstance(m, torch.nn.Embedding),
+        )
+        ep = torch.export.export(model, example_inputs)
+        delegated_program = executorch.exir.to_edge_transform_and_lower(
+            ep,
+            partitioner=[self._coreml_partitioner()],
+        )
+        for node in delegated_program.exported_program().graph.nodes:
+            if node.op == "call_function":
+                assert node.target.__name__ in [
+                    "executorch_call_delegate",
+                    "getitem",
+                ], f"Got unexpected node target after delegation: {node.target.__name__}"
+        et_prog = delegated_program.to_executorch()
+        print(et_prog.exported_program())
+        self._compare_outputs(et_prog, model, example_inputs)
+
 
 if __name__ == "__main__":
     test_runner = TestTorchOps()
-    # test_runner.test_dequantize_affine_b4w_embedding()
-    # test_runner.test_dequantize_affine_b4w_linear()
-    # test_runner.test_dequantize_affine_c4w_embedding()
-    # test_runner.test_dequantize_affine_c4w_linear()
-    # test_runner.test_dequantize_affine_c8w_embedding_b4w_linear()
+    test_runner.test_dequantize_affine_b4w_embedding()
+    test_runner.test_dequantize_affine_b4w_linear()
+    test_runner.test_dequantize_affine_c4w_embedding()
+    test_runner.test_dequantize_affine_c4w_linear()
+    test_runner.test_dequantize_affine_c8w_embedding_b4w_linear()
     test_runner.test_dequantize_codebook_linear()
+    test_runner.test_dequantize_codebook_embedding()