From 7de3c9f41c4ff7c200802fc63fb5f38efe46ed82 Mon Sep 17 00:00:00 2001 From: usimha <135899523+u-simha@users.noreply.github.com> Date: Fri, 28 Aug 2026 16:15:23 -0700 Subject: [PATCH 1/6] Add joint post-training quantization/palettization + sparsity support Adds a hidden, settable-but-not-public `_sparsity` field to QuantizationSpec and PalettizationSpec: weights are pre-sparsified before fake-quant/palettize, and finalize() inserts sparse_to_dense (plus lut_to_dense for palettization) in the correct op order. Validation of unsupported combinations (asymmetric quant, unsigned/FP4 dtypes, quantized LUTs, per-channel scale, non-per-tensor granularity) now happens at spec-construction time via pydantic. --- .../kmeans/_prepare_for_export.py | 49 +++++++- .../palettization/kmeans/palettizer.py | 1 + .../palettization/spec/fake_palettize.py | 10 ++ src/coreai_opt/palettization/spec/spec.py | 21 ++++ .../_graph/_prepare_for_export.py | 31 ++++- src/coreai_opt/quantization/spec/factory.py | 2 + .../quantization/spec/fake_quantize.py | 10 ++ src/coreai_opt/quantization/spec/spec.py | 26 ++++ tests/export/test_joint_sparsity.py | 119 ++++++++++++++++++ 9 files changed, 262 insertions(+), 7 deletions(-) create mode 100644 tests/export/test_joint_sparsity.py diff --git a/src/coreai_opt/palettization/kmeans/_prepare_for_export.py b/src/coreai_opt/palettization/kmeans/_prepare_for_export.py index 678188da..81b9c749 100644 --- a/src/coreai_opt/palettization/kmeans/_prepare_for_export.py +++ b/src/coreai_opt/palettization/kmeans/_prepare_for_export.py @@ -6,6 +6,7 @@ from dataclasses import dataclass from os import PathLike from pathlib import Path +from typing import Any import torch import torch.nn as nn @@ -46,6 +47,32 @@ class PalettizationInfo: lut_quantization: LUTQuantizationInfo | None = None +class _SparsePalettizeReconstruction(nn.Module): + """Parametrization module inserted to reconstruct a sparse-palettized weight. + + Traces ``coreai.lut_to_dense`` and ``coreai.sparse_to_dense`` in that order. + """ + + def __init__( + self, + nonzero_indices: torch.Tensor, + lut: torch.Tensor, + mask: torch.Tensor, + vector_axis: int | None, + ) -> None: + super().__init__() + self.register_buffer("nonzero_indices", nonzero_indices) + self.register_buffer("lut", lut.reshape(1, lut.shape[-2], lut.shape[-1])) + self.register_buffer("mask", mask) + self.vector_axis = 0 if vector_axis is None else vector_axis + + def forward(self, _: Any) -> torch.Tensor: + nonzero_values = torch.ops.coreai.lut_to_dense( + self.nonzero_indices, self.lut, self.vector_axis + ) + return torch.ops.coreai.sparse_to_dense(nonzero_values, self.mask) + + def _expand_rank( tensor: torch.Tensor, target_rank: int, @@ -210,6 +237,7 @@ def _insert_mlir_custom_op( module_name: str, param_name: str, palett_info: PalettizationInfo, + fake_palett_mod: _FakePalettizeImplBase, fake_palett_idx: int, mmap_dir: str | PathLike[str] | None, ) -> None: @@ -227,6 +255,10 @@ def _insert_mlir_custom_op( 4. Both: lut_to_dense(int LUT) + constexpr_blockwise_shift_scale(fused_scale) where fused_scale = lut_scale * per_channel_scale + When ``fake_palett_mod.sparsity`` is set, the LUT lookup runs on the + nonzero-only indices and the result is packed via ``coreai::sparse_to_dense`` + instead of installing a plain Palettize/ScaledPalettize parametrization. + When ``mmap_dir`` is provided, the new MLIR module is serialized to a safetensors file under that directory and reloaded via mmap before being swapped in. @@ -261,7 +293,20 @@ def _import_coreai_torch_modules(): vector_axis = _DEFAULT_VECTOR_AXIS if palett_info.cluster_dim > 1 else None - if needs_scale: + if fake_palett_mod.sparsity is not None: + # Reuses the mask from prepare()'s forward pass. needs_scale is always + # False here: PalettizationSpec rejects lut_qspec/enable_per_channel_scale + # combined with sparsity, since both are position-dependent and would + # be scrambled by flattening to the nonzero-only indices below. + mask = fake_palett_mod._sparsity_mask.to(torch.bool) + nonzero_indices = palett_info.indices[mask] + mlir_palett_mod = _SparsePalettizeReconstruction( + nonzero_indices=nonzero_indices, + lut=palett_info.lut, + mask=mask, + vector_axis=vector_axis, + ) + elif needs_scale: lut, scale, zero_point = _resolve_mlir_lut_and_scale(palett_info) mlir_palett_mod = ScaledPalettizeParametrization( indices=palett_info.indices, @@ -336,7 +381,7 @@ def _process_palettized_parameter( _register_mil_compression_metadata(module, param_name, palett_info) elif backend == ExportBackend.CoreAI: _insert_mlir_custom_op( - module, module_name, param_name, palett_info, fake_palett_idx, mmap_dir + module, module_name, param_name, palett_info, fake_palett_mod, fake_palett_idx, mmap_dir ) diff --git a/src/coreai_opt/palettization/kmeans/palettizer.py b/src/coreai_opt/palettization/kmeans/palettizer.py index dee63ab7..b36e5298 100644 --- a/src/coreai_opt/palettization/kmeans/palettizer.py +++ b/src/coreai_opt/palettization/kmeans/palettizer.py @@ -432,6 +432,7 @@ def _spec_to_partial( # Serialize the spec, then layer in the owning module's compressor-specific # settings (e.g. enable_fast_kmeans_mode, rounding_precision). args = spec.model_dump_preserve_objects() + args["sparsity"] = spec._sparsity args.update(module_config._get_compressor_specific_settings()) return _KMeansFakePalettize.with_args(**args) diff --git a/src/coreai_opt/palettization/spec/fake_palettize.py b/src/coreai_opt/palettization/spec/fake_palettize.py index ee4f7e36..c5538d81 100644 --- a/src/coreai_opt/palettization/spec/fake_palettize.py +++ b/src/coreai_opt/palettization/spec/fake_palettize.py @@ -23,6 +23,7 @@ _IncompatibleClusterDimError, _IncompatibleGranularityError, ) +from coreai_opt.pruning.spec import PruneImplBase, Unstructured from coreai_opt.quantization.spec import QuantizationSpec logger = logging.getLogger(__name__) @@ -49,6 +50,7 @@ def __init__( granularity: PalettizationGranularity, cluster_dim: int, enable_per_channel_scale: bool, + sparsity: float | None = None, **kwargs, ): super().__init__(**kwargs) @@ -57,10 +59,12 @@ def __init__( self.granularity = granularity self.cluster_dim = cluster_dim self.enable_per_channel_scale = enable_per_channel_scale + self.sparsity = sparsity self.register_buffer("fake_palett_enabled", torch.tensor([1], dtype=torch.uint8)) self.register_buffer("observer_enabled", torch.tensor([1], dtype=torch.uint8)) self._disabled = False + self.register_buffer("_sparsity_mask", None, persistent=False) self.register_buffer("lut", None) self.register_buffer("indices", None) @@ -84,6 +88,12 @@ def forward(self, tensor: torch.Tensor) -> torch.Tensor: if self._disabled: return tensor + if self.sparsity is not None: + self._sparsity_mask = PruneImplBase.resolve("default").compute_mask( + tensor, self.sparsity, Unstructured() + ) + tensor = tensor * self._sparsity_mask + if self.observer_enabled[0] == 1: # Cluster weights try: diff --git a/src/coreai_opt/palettization/spec/spec.py b/src/coreai_opt/palettization/spec/spec.py index d3b47348..d095c386 100644 --- a/src/coreai_opt/palettization/spec/spec.py +++ b/src/coreai_opt/palettization/spec/spec.py @@ -95,6 +95,27 @@ class PalettizationSpec(CompressionSpec): # Private attribute for compression type _compression_type: CompressionType = PrivateAttr(default=CompressionType.PALETTIZATION) + # Sparsity level, in [0, 1]. Set via the `_sparsity` constructor/dict key. + _sparsity: float | None = PrivateAttr(default=None) + + def __init__(self, **data: Any) -> None: + sparsity = data.pop("_sparsity", None) + super().__init__(**data) + if sparsity is not None: + if not (0.0 <= sparsity <= 1.0): + raise ValueError(f"_sparsity must be in [0, 1], got {sparsity}") + self._validate_sparsity(sparsity) + self._sparsity = sparsity + + def _validate_sparsity(self, sparsity: float) -> None: + """Reject sparsity combined with a position-dependent LUT/scale mapping.""" + if self.lut_qspec is not None: + raise ValueError("lut_qspec not supported for joint sparsity.") + if not isinstance(self.granularity, PerTensorGranularity): + raise ValueError(f"granularity={self.granularity} not supported for joint sparsity.") + if self.enable_per_channel_scale: + raise ValueError("enable_per_channel_scale not supported for joint sparsity.") + @model_validator(mode="after") def validate_lut_qspec(self) -> "PalettizationSpec": """Validate that lut_qspec only uses supported configurations.""" diff --git a/src/coreai_opt/quantization/_graph/_prepare_for_export.py b/src/coreai_opt/quantization/_graph/_prepare_for_export.py index 60f15459..55e6bda2 100644 --- a/src/coreai_opt/quantization/_graph/_prepare_for_export.py +++ b/src/coreai_opt/quantization/_graph/_prepare_for_export.py @@ -272,14 +272,20 @@ def _import_coreai_custom_ops(): # Extract and prepare quantization parameters scale, zero_point, minval = _extract_quantization_params(fake_quant_mod) + # Cast scale and minval to appropriate dtype for MLIR backend inference _compute_dtype_for_export = fake_quant_mod.qparams_calculator._compute_dtype_for_export scale = scale.to(dtype=_compute_dtype_for_export) if minval is not None: minval = minval.to(dtype=_compute_dtype_for_export) - # Construct quantized weights + # Construct quantized weights, reusing the mask computed during + # prepare()'s forward pass if sparsity is set. dense_weight = resolve_attr(model, input_node.target).data + mask: torch.Tensor | None = None + if fake_quant_mod.sparsity is not None: + mask = fake_quant_mod._sparsity_mask.to(torch.bool) + dense_weight = dense_weight * mask quantized_data = fake_quant_mod.quantize(dense_weight, scale, zero_point, minval) # Drop one of the offsets so that the export @@ -295,13 +301,28 @@ def _import_coreai_custom_ops(): # Register buffers and get buffer names param_name = str(input_node.target).replace(".", "_") - buffer_names = _register_quantization_buffers( - model, param_name, scale, zero_point, quantized_data, minval - ) + if mask is not None: + nonzero_data = quantized_data[mask] + buffer_names = _register_quantization_buffers( + model, param_name, scale, zero_point, minval=minval + ) + model.register_buffer(f"{param_name}_nonzero", nonzero_data) + model.register_buffer(f"{param_name}_mask", mask) + else: + buffer_names = _register_quantization_buffers( + model, param_name, scale, zero_point, quantized_data, minval + ) # Create graph nodes and replace fake quantization with model.graph.inserting_before(node): - quantized_data_node = model.graph.get_attr(buffer_names["quantized_data"]) + if mask is not None: + nonzero_node = model.graph.get_attr(f"{param_name}_nonzero") + mask_node = model.graph.get_attr(f"{param_name}_mask") + quantized_data_node = model.graph.call_function( + coreai.sparse_to_dense, (nonzero_node, mask_node) + ) + else: + quantized_data_node = model.graph.get_attr(buffer_names["quantized_data"]) scale_node = model.graph.get_attr(buffer_names["scale"]) if zero_point is not None: diff --git a/src/coreai_opt/quantization/spec/factory.py b/src/coreai_opt/quantization/spec/factory.py index 9a449b0d..928f9448 100644 --- a/src/coreai_opt/quantization/spec/factory.py +++ b/src/coreai_opt/quantization/spec/factory.py @@ -203,6 +203,7 @@ def create_fake_quantizer( "qparams_calculator": qparams_calculator, "quantization_target": quantization_target, "n_bits": spec.n_bits, + "sparsity": spec._sparsity, } # Automatically detect and include any extra arguments @@ -244,6 +245,7 @@ def create_fake_quantizer_partial( "quant_max": spec.quant_max, "quantization_target": quantization_target, "n_bits": spec.n_bits, + "sparsity": spec._sparsity, } # Automatically detect and include any extra arguments diff --git a/src/coreai_opt/quantization/spec/fake_quantize.py b/src/coreai_opt/quantization/spec/fake_quantize.py index aafd4768..5d50b1e3 100644 --- a/src/coreai_opt/quantization/spec/fake_quantize.py +++ b/src/coreai_opt/quantization/spec/fake_quantize.py @@ -26,6 +26,7 @@ is_float_quant_dtype as _is_float_quant_dtype, ) from coreai_opt.config.spec import CompressionSimulatorBase, CompressionTargetTensor +from coreai_opt.pruning.spec import PruneImplBase, Unstructured from coreai_opt.quantization._utils import get_quantization_shapes as _get_quantization_shapes from coreai_opt.quantization.spec.errors import _BlockSizeMismatchError @@ -56,6 +57,7 @@ def __init__( qparams_calculator: QParamsCalculatorBase, quantization_target: CompressionTargetTensor, n_bits: int | None = None, + sparsity: float | None = None, **kwargs, ): super().__init__() @@ -68,7 +70,9 @@ def __init__( self.quant_max = quant_max self.qparams_calculator = qparams_calculator self.quantization_target = quantization_target + self.sparsity = sparsity self.register_buffer("_disabled", torch.tensor(False)) + self.register_buffer("_sparsity_mask", None, persistent=False) # Infer n_bits from dtype if not provided if n_bits is None: @@ -140,6 +144,12 @@ def forward(self, tensor: torch.Tensor) -> torch.Tensor: if self._disabled.item(): return tensor + if self.sparsity is not None: + self._sparsity_mask = PruneImplBase.resolve("default").compute_mask( + tensor, self.sparsity, Unstructured() + ) + tensor = tensor * self._sparsity_mask + if self.observer_enabled[0] == 1: # Call the forward function of the qparams_calculator # to collect observer statistics when the observer is diff --git a/src/coreai_opt/quantization/spec/spec.py b/src/coreai_opt/quantization/spec/spec.py index 52a83722..b724fd17 100644 --- a/src/coreai_opt/quantization/spec/spec.py +++ b/src/coreai_opt/quantization/spec/spec.py @@ -372,6 +372,18 @@ class type: MinMaxRangeCalculator or custom registered class type # Private attribute for compression type _compression_type: CompressionType = PrivateAttr(default=CompressionType.QUANTIZATION) + # Sparsity level, in [0, 1]. Set via the `_sparsity` constructor/dict key. + _sparsity: float | None = PrivateAttr(default=None) + + def __init__(self, **data: Any) -> None: + sparsity = data.pop("_sparsity", None) + super().__init__(**data) + if sparsity is not None: + if not (0.0 <= sparsity <= 1.0): + raise ValueError(f"_sparsity must be in [0, 1], got {sparsity}") + self._validate_sparsity_zero_preserving(sparsity) + self._sparsity = sparsity + # Supported dtypes for quantization (class attribute for testing extensibility) SUPPORTED_DTYPES: ClassVar[set[torch.dtype]] = { # Signed integer types @@ -555,6 +567,20 @@ def validate_scale_dtype(self) -> QuantizationSpec: return self + def _validate_sparsity_zero_preserving(self, sparsity: float) -> None: + """Reject sparsity unless a raw 0 dequantizes to exactly 0.0.""" + if _is_float4_dtype(self.dtype): + raise ValueError("FP4 dtype not supported for joint sparsity.") + if self.dtype.is_floating_point: + return + + if self.qformulation != QuantizationFormulation.ZP: + raise ValueError(f"qformulation={self.qformulation} not supported for joint sparsity.") + if self.qscheme == QuantizationScheme.ASYMMETRIC: + raise ValueError(f"qscheme={self.qscheme} not supported for joint sparsity.") + if not self.dtype.is_signed: + raise ValueError(f"unsigned dtype={self.dtype} not supported for joint sparsity.") + def get_extra_args(self) -> dict[str, Any]: """ Automatically detect and return fields beyond base QuantizationSpec. diff --git a/tests/export/test_joint_sparsity.py b/tests/export/test_joint_sparsity.py new file mode 100644 index 00000000..5957d597 --- /dev/null +++ b/tests/export/test_joint_sparsity.py @@ -0,0 +1,119 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""End-to-end export tests for joint post-training quantization/palettization + sparsity.""" + +import pytest +import torch +import torch.nn as nn + +from coreai_opt import ExportBackend +from coreai_opt.palettization import ( + KMeansPalettizer, + KMeansPalettizerConfig, + ModuleKMeansPalettizerConfig, + PalettizationSpec, +) +from coreai_opt.quantization import ModuleQuantizerConfig, Quantizer, QuantizerConfig +from coreai_opt.quantization.config import ExecutionMode +from coreai_opt.quantization.spec import PerTensorGranularity, QuantizationScheme, QuantizationSpec + +from . import export_utils + + +class TestJointSparsityExport: + """PTQ/PTP + PTS (post-training quantization/palettization + sparsity), end to end.""" + + @staticmethod + def _run_quant_sparsity_export( + model: nn.Module, input_data: torch.Tensor, expected_count: int + ) -> None: + model.eval() + config = QuantizerConfig( + global_config=ModuleQuantizerConfig( + op_state_spec={ + "weight": QuantizationSpec( + dtype=torch.int8, + qscheme=QuantizationScheme.SYMMETRIC, + granularity=PerTensorGranularity(), + _sparsity=0.5, + ) + }, + op_input_spec=None, + op_output_spec=None, + ), + execution_mode=ExecutionMode.GRAPH, + ) + + quantizer = Quantizer(model, config) + prepared_model = quantizer.prepare((input_data,)) + + with torch.no_grad(): + prepared_model_output = prepared_model(input_data) + + finalized_model = quantizer.finalize(backend=ExportBackend.CoreAI) + + export_utils.convert_and_verify( + finalized_model=finalized_model, + input_data=input_data, + expected_ops={ + "sparse_to_dense": expected_count, + "constexpr_blockwise_shift_scale": expected_count, + }, + export_backend=ExportBackend.CoreAI, + prepared_model_output=prepared_model_output, + ) + + @staticmethod + def _run_palettization_sparsity_export( + model: nn.Module, input_data: torch.Tensor, expected_count: int + ) -> None: + model.eval() + config = KMeansPalettizerConfig( + global_config=ModuleKMeansPalettizerConfig( + op_state_spec={"weight": PalettizationSpec(n_bits=8, _sparsity=0.5)} + ) + ) + + palettizer = KMeansPalettizer(model, config) + prepared_model = palettizer.prepare((input_data,)) + + with torch.no_grad(): + prepared_model_output = prepared_model(input_data) + + finalized_model = palettizer.finalize(backend=ExportBackend.CoreAI) + + export_utils.convert_and_verify( + finalized_model=finalized_model, + input_data=input_data, + expected_ops={ + "lut_to_dense": expected_count, + "sparse_to_dense": expected_count, + }, + export_backend=ExportBackend.CoreAI, + prepared_model_output=prepared_model_output, + ) + + def test_quant_sparsity_mnist_export(self, custom_test_mnist_model, mnist_example_input): + self._run_quant_sparsity_export( + custom_test_mnist_model, mnist_example_input, expected_count=6 + ) + + @pytest.mark.slow + def test_quant_sparsity_resnet_export(self, resnet50_model, resnet_example_input): + self._run_quant_sparsity_export(resnet50_model, resnet_example_input, expected_count=54) + + def test_palettization_sparsity_mnist_export( + self, custom_test_mnist_model, mnist_example_input + ): + self._run_palettization_sparsity_export( + custom_test_mnist_model, mnist_example_input, expected_count=6 + ) + + @pytest.mark.slow + def test_palettization_sparsity_resnet_export(self, resnet50_model, resnet_example_input): + self._run_palettization_sparsity_export( + resnet50_model, resnet_example_input, expected_count=54 + ) From 7ce24def1431d2db51d9c9f947b5b88a0ccec06e Mon Sep 17 00:00:00 2001 From: usimha <135899523+u-simha@users.noreply.github.com> Date: Mon, 31 Aug 2026 14:43:45 -0700 Subject: [PATCH 2/6] Trim sparsity docstrings/comments and consolidate range validation Moves the [0, 1] range check for _sparsity into the existing _validate_sparsity(_zero_preserving) methods so each spec has a single validation entrypoint, and shortens a couple of over-long comments. --- .../palettization/kmeans/_prepare_for_export.py | 8 ++++---- src/coreai_opt/palettization/spec/spec.py | 4 ++-- src/coreai_opt/quantization/spec/spec.py | 4 ++-- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/src/coreai_opt/palettization/kmeans/_prepare_for_export.py b/src/coreai_opt/palettization/kmeans/_prepare_for_export.py index 81b9c749..20bc2f7d 100644 --- a/src/coreai_opt/palettization/kmeans/_prepare_for_export.py +++ b/src/coreai_opt/palettization/kmeans/_prepare_for_export.py @@ -62,6 +62,8 @@ def __init__( ) -> None: super().__init__() self.register_buffer("nonzero_indices", nonzero_indices) + # nonzero_indices is rank 1 (flattened by masking), so lut must be + # reshaped to rank 3 (lut_to_dense requires lut.rank == indices.rank + 2). self.register_buffer("lut", lut.reshape(1, lut.shape[-2], lut.shape[-1])) self.register_buffer("mask", mask) self.vector_axis = 0 if vector_axis is None else vector_axis @@ -294,10 +296,8 @@ def _import_coreai_torch_modules(): vector_axis = _DEFAULT_VECTOR_AXIS if palett_info.cluster_dim > 1 else None if fake_palett_mod.sparsity is not None: - # Reuses the mask from prepare()'s forward pass. needs_scale is always - # False here: PalettizationSpec rejects lut_qspec/enable_per_channel_scale - # combined with sparsity, since both are position-dependent and would - # be scrambled by flattening to the nonzero-only indices below. + # Reuses the mask from prepare()'s forward pass. needs_scale is always False here: + # PalettizationSpec rejects lut_qspec/enable_per_channel_scale combined with sparsity. mask = fake_palett_mod._sparsity_mask.to(torch.bool) nonzero_indices = palett_info.indices[mask] mlir_palett_mod = _SparsePalettizeReconstruction( diff --git a/src/coreai_opt/palettization/spec/spec.py b/src/coreai_opt/palettization/spec/spec.py index d095c386..798c0f73 100644 --- a/src/coreai_opt/palettization/spec/spec.py +++ b/src/coreai_opt/palettization/spec/spec.py @@ -102,13 +102,13 @@ def __init__(self, **data: Any) -> None: sparsity = data.pop("_sparsity", None) super().__init__(**data) if sparsity is not None: - if not (0.0 <= sparsity <= 1.0): - raise ValueError(f"_sparsity must be in [0, 1], got {sparsity}") self._validate_sparsity(sparsity) self._sparsity = sparsity def _validate_sparsity(self, sparsity: float) -> None: """Reject sparsity combined with a position-dependent LUT/scale mapping.""" + if not (0.0 <= sparsity <= 1.0): + raise ValueError(f"_sparsity must be in [0, 1], got {sparsity}") if self.lut_qspec is not None: raise ValueError("lut_qspec not supported for joint sparsity.") if not isinstance(self.granularity, PerTensorGranularity): diff --git a/src/coreai_opt/quantization/spec/spec.py b/src/coreai_opt/quantization/spec/spec.py index b724fd17..e776c218 100644 --- a/src/coreai_opt/quantization/spec/spec.py +++ b/src/coreai_opt/quantization/spec/spec.py @@ -379,8 +379,6 @@ def __init__(self, **data: Any) -> None: sparsity = data.pop("_sparsity", None) super().__init__(**data) if sparsity is not None: - if not (0.0 <= sparsity <= 1.0): - raise ValueError(f"_sparsity must be in [0, 1], got {sparsity}") self._validate_sparsity_zero_preserving(sparsity) self._sparsity = sparsity @@ -569,6 +567,8 @@ def validate_scale_dtype(self) -> QuantizationSpec: def _validate_sparsity_zero_preserving(self, sparsity: float) -> None: """Reject sparsity unless a raw 0 dequantizes to exactly 0.0.""" + if not (0.0 <= sparsity <= 1.0): + raise ValueError(f"_sparsity must be in [0, 1], got {sparsity}") if _is_float4_dtype(self.dtype): raise ValueError("FP4 dtype not supported for joint sparsity.") if self.dtype.is_floating_point: From 19d640b9fe6b8302290f7c151260cc33b0bf1ec3 Mon Sep 17 00:00:00 2001 From: usimha <135899523+u-simha@users.noreply.github.com> Date: Fri, 11 Sep 2026 15:47:52 -0700 Subject: [PATCH 3/6] Move joint-sparsity restrictions from spec to export, wire CoreML sparse ops Spec construction no longer validates sparsity-compatibility; instead _validate_sparsity_for_export (in each technique's _prepare_for_export.py) checks at finalize time: quantization requires a zero (or absent) zero_point, palettization requires per-tensor granularity and scalar (cluster_dim=1) palettization. The same check now gates both backends -- CoreML's export path registers a PRUNING-first compression_type list so coremltools' own torch-frontend converter auto-detects the sparsity pattern from the traced weight's real zeros and builds the matching constexpr_sparse_to_dense / constexpr_sparse_blockwise_shift_scale / constexpr_lut_to_sparse chain, mirroring CoreAI's sparse_to_dense / lut_to_dense chain. Adds tests/export/test_joint_compression.py, covering the quantization dtype/qscheme/granularity matrix and the palettization n_bits/cluster_dim/ granularity matrix against both export backends, on MNIST and ResNet. --- .../kmeans/_prepare_for_export.py | 45 ++- src/coreai_opt/palettization/spec/spec.py | 14 +- .../_graph/_prepare_for_export.py | 30 +- src/coreai_opt/quantization/spec/spec.py | 19 +- tests/export/test_joint_compression.py | 323 ++++++++++++++++++ 5 files changed, 396 insertions(+), 35 deletions(-) create mode 100644 tests/export/test_joint_compression.py diff --git a/src/coreai_opt/palettization/kmeans/_prepare_for_export.py b/src/coreai_opt/palettization/kmeans/_prepare_for_export.py index a6172ff0..81e453da 100644 --- a/src/coreai_opt/palettization/kmeans/_prepare_for_export.py +++ b/src/coreai_opt/palettization/kmeans/_prepare_for_export.py @@ -24,6 +24,7 @@ from coreai_opt.palettization.spec.fake_palettize import ( _FakePalettizeImplBase, ) +from coreai_opt.palettization.spec.granularity import PerTensorGranularity _DEFAULT_VECTOR_AXIS = 0 @@ -145,6 +146,7 @@ def _register_mil_compression_metadata( module: nn.Module, param_name: str, palett_info: PalettizationInfo, + fake_palett_mod: _FakePalettizeImplBase, ) -> None: """ Remove the fake palettization parametrization from the module @@ -161,12 +163,20 @@ def _register_mil_compression_metadata( leave_parametrized=True, ) - # Determine compression type(s) + # Determine compression type(s). coremltools' own torch-frontend converter + # auto-detects sparsity from raw zeros in the traced weight value when + # compression_type lists PRUNING first, then chains PALETTIZATION onto its + # constexpr_sparse_to_dense output (constexpr_lut_to_sparse) -- which has the + # same flatten-to-nonzero-only-indices contract as our own CoreAI reconstruction, + # so it needs the same per-tensor/scalar guarantee. lut_quant = palett_info.lut_quantization if lut_quant is not None: compression_type = [CompressionType.PALETTIZATION, CompressionType.QUANTIZATION] else: - compression_type = CompressionType.PALETTIZATION + compression_type = [CompressionType.PALETTIZATION] + if fake_palett_mod.sparsity is not None: + _validate_sparsity_for_export(fake_palett_mod) + compression_type = [CompressionType.PRUNING, *compression_type] metadata = MILCompressionMetadata( param_name=param_name, @@ -232,6 +242,31 @@ def _resolve_mlir_lut_and_scale( return lut, scale, offset +def _validate_sparsity_for_export(fake_palett_mod: _FakePalettizeImplBase) -> None: + """Reject sparsity combined with anything that isn't a single, position-independent LUT. + + Masking flattens indices to rank 1 before the LUT lookup, which only preserves + meaning when every element shares one scalar codebook: per-tensor granularity + (not per-channel/grouped, which use multiple LUTs) and cluster_dim == 1 (not + vector palettization, whose indices are already at a reduced, position-dependent + rank). A quantized LUT or per-channel scale would also apply a position-dependent + mapping that flattening would scramble. + """ + if not isinstance(fake_palett_mod.granularity, PerTensorGranularity): + raise ValueError( + f"granularity={fake_palett_mod.granularity} not supported for joint sparsity." + ) + if fake_palett_mod.cluster_dim != 1: + raise ValueError( + f"cluster_dim={fake_palett_mod.cluster_dim} (vector palettization) " + "not supported for joint sparsity." + ) + if fake_palett_mod.lut_qspec is not None: + raise ValueError("lut_qspec not supported for joint sparsity.") + if fake_palett_mod.enable_per_channel_scale: + raise ValueError("enable_per_channel_scale not supported for joint sparsity.") + + def _insert_mlir_custom_op( module: nn.Module, module_name: str, @@ -294,8 +329,8 @@ def _import_coreai_torch_modules(): vector_axis = _DEFAULT_VECTOR_AXIS if palett_info.cluster_dim > 1 else None if fake_palett_mod.sparsity is not None: - # Reuses the mask from prepare()'s forward pass. needs_scale is always False here: - # PalettizationSpec rejects lut_qspec/enable_per_channel_scale combined with sparsity. + _validate_sparsity_for_export(fake_palett_mod) + # Reuses the mask from prepare()'s forward pass. mask = fake_palett_mod._sparsity_mask.to(torch.bool) nonzero_indices = palett_info.indices[mask] mlir_palett_mod = _SparsePalettizeReconstruction( @@ -376,7 +411,7 @@ def _process_palettized_parameter( ) if backend == ExportBackend.CoreML: - _register_mil_compression_metadata(module, param_name, palett_info) + _register_mil_compression_metadata(module, param_name, palett_info, fake_palett_mod) elif backend == ExportBackend.CoreAI: _insert_mlir_custom_op( module, module_name, param_name, palett_info, fake_palett_mod, fake_palett_idx, mmap_dir diff --git a/src/coreai_opt/palettization/spec/spec.py b/src/coreai_opt/palettization/spec/spec.py index 95749157..7148ca8c 100644 --- a/src/coreai_opt/palettization/spec/spec.py +++ b/src/coreai_opt/palettization/spec/spec.py @@ -115,20 +115,10 @@ def __init__(self, **data: Any) -> None: sparsity = data.pop("_sparsity", None) super().__init__(**data) if sparsity is not None: - self._validate_sparsity(sparsity) + if not (0.0 <= sparsity <= 1.0): + raise ValueError(f"_sparsity must be in [0, 1], got {sparsity}") self._sparsity = sparsity - def _validate_sparsity(self, sparsity: float) -> None: - """Reject sparsity combined with a position-dependent LUT/scale mapping.""" - if not (0.0 <= sparsity <= 1.0): - raise ValueError(f"_sparsity must be in [0, 1], got {sparsity}") - if self.lut_qspec is not None: - raise ValueError("lut_qspec not supported for joint sparsity.") - if not isinstance(self.granularity, PerTensorGranularity): - raise ValueError(f"granularity={self.granularity} not supported for joint sparsity.") - if self.enable_per_channel_scale: - raise ValueError("enable_per_channel_scale not supported for joint sparsity.") - @model_validator(mode="after") def validate_lut_qspec(self) -> "PalettizationSpec": """Validate that lut_qspec only uses supported configurations.""" diff --git a/src/coreai_opt/quantization/_graph/_prepare_for_export.py b/src/coreai_opt/quantization/_graph/_prepare_for_export.py index 9da0e660..524f0b8a 100644 --- a/src/coreai_opt/quantization/_graph/_prepare_for_export.py +++ b/src/coreai_opt/quantization/_graph/_prepare_for_export.py @@ -230,6 +230,22 @@ def _register_quantization_buffers( return buffer_names +def _validate_sparsity_for_export( + fake_quant_mod: FakeQuantizeImplBase, zero_point: torch.Tensor | None +) -> None: + """Reject sparsity combined with anything ``sparse_to_dense``'s raw-0 padding can't represent. + + ``sparse_to_dense`` always pads pruned positions with a raw literal 0, so joint + sparsity is only correct when that 0 dequantizes to exactly 0.0, i.e. ``zero_point`` + is 0 (or absent). FP4 is rejected separately: its packing runs before the nonzero + values are extracted, and would misalign against the full-resolution mask. + """ + if is_float4_dtype(fake_quant_mod.dtype): + raise ValueError("FP4 dtype not supported for joint sparsity.") + if zero_point is not None and not torch.all(zero_point == 0): + raise ValueError("nonzero zero_point not supported for joint sparsity.") + + def _process_mlir_weight_quantization( model: torch.fx.GraphModule, node: Node, @@ -272,6 +288,7 @@ def _import_coreai_custom_ops(): dense_weight = resolve_attr(model, input_node.target).data mask: torch.Tensor | None = None if fake_quant_mod.sparsity is not None: + _validate_sparsity_for_export(fake_quant_mod, zero_point) mask = fake_quant_mod._sparsity_mask.to(torch.bool) dense_weight = dense_weight * mask quantized_data = fake_quant_mod.quantize(dense_weight, scale, zero_point, minval) @@ -481,10 +498,21 @@ def _process_mil_weight_quantization( # Get the module that owns the weight parameter weight_module: torch.nn.Module = _get_weight_module(modules, module_name) + # coremltools' own torch-frontend converter auto-detects sparsity from raw + # zeros in the registered weight value when compression_type lists PRUNING + # first, then chains QUANTIZATION onto its constexpr_sparse_to_dense output + # (see coremltools.converters.mil.frontend.torch.converter._construct_compression_op). + # That chain has the same raw-0-pads-pruned-positions contract as our own + # sparse_to_dense, so it needs the same zero-preserving guarantee. + compression_type = [CompressionType.QUANTIZATION] + if fake_quant_mod.sparsity is not None: + _validate_sparsity_for_export(fake_quant_mod, zero_point) + compression_type = [CompressionType.PRUNING, CompressionType.QUANTIZATION] + # Create and register metadata metadata = MILCompressionMetadata( param_name=param_name, - compression_type=CompressionType.QUANTIZATION, + compression_type=compression_type, quantization_n_bits=fake_quant_mod.n_bits, quantization_scale=scale, zero_point=zero_point, diff --git a/src/coreai_opt/quantization/spec/spec.py b/src/coreai_opt/quantization/spec/spec.py index 0495c2f2..afe4c9da 100644 --- a/src/coreai_opt/quantization/spec/spec.py +++ b/src/coreai_opt/quantization/spec/spec.py @@ -380,7 +380,8 @@ def __init__(self, **data: Any) -> None: sparsity = data.pop("_sparsity", None) super().__init__(**data) if sparsity is not None: - self._validate_sparsity_zero_preserving(sparsity) + if not (0.0 <= sparsity <= 1.0): + raise ValueError(f"_sparsity must be in [0, 1], got {sparsity}") self._sparsity = sparsity # Supported dtypes for quantization (class attribute for testing extensibility) @@ -566,22 +567,6 @@ def validate_scale_dtype(self) -> QuantizationSpec: return self - def _validate_sparsity_zero_preserving(self, sparsity: float) -> None: - """Reject sparsity unless a raw 0 dequantizes to exactly 0.0.""" - if not (0.0 <= sparsity <= 1.0): - raise ValueError(f"_sparsity must be in [0, 1], got {sparsity}") - if _is_float4_dtype(self.dtype): - raise ValueError("FP4 dtype not supported for joint sparsity.") - if self.dtype.is_floating_point: - return - - if self.qformulation != QuantizationFormulation.ZP: - raise ValueError(f"qformulation={self.qformulation} not supported for joint sparsity.") - if self.qscheme == QuantizationScheme.ASYMMETRIC: - raise ValueError(f"qscheme={self.qscheme} not supported for joint sparsity.") - if not self.dtype.is_signed: - raise ValueError(f"unsigned dtype={self.dtype} not supported for joint sparsity.") - def get_extra_args(self) -> dict[str, Any]: """ Automatically detect and return fields beyond base QuantizationSpec. diff --git a/tests/export/test_joint_compression.py b/tests/export/test_joint_compression.py new file mode 100644 index 00000000..13805657 --- /dev/null +++ b/tests/export/test_joint_compression.py @@ -0,0 +1,323 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +import pytest +import torch +import torch.nn as nn + +from coreai_opt import ExportBackend +from coreai_opt.palettization import ( + KMeansPalettizer, + KMeansPalettizerConfig, + ModuleKMeansPalettizerConfig, + PalettizationSpec, +) +from coreai_opt.palettization.spec import ( + PerGroupedChannelGranularity, + PerTensorGranularity as PalettPerTensorGranularity, +) +from coreai_opt.quantization import ModuleQuantizerConfig, Quantizer, QuantizerConfig +from coreai_opt.quantization.config import ExecutionMode +from coreai_opt.quantization.spec import ( + PerChannelGranularity, + PerTensorGranularity, + QuantizationGranularity, + QuantizationScheme, + QuantizationSpec, +) + +from . import export_utils + +_SPARSITY = 0.5 +_MNIST_LAYER_COUNT = 6 +_RESNET_LAYER_COUNT = 54 +_BACKENDS = [ExportBackend.CoreAI, ExportBackend.CoreML] + +_QUANT_EXPECTED_OPS = { + ExportBackend.CoreAI: lambda n: { + "sparse_to_dense": n, + "constexpr_blockwise_shift_scale": n, + }, + ExportBackend.CoreML: lambda n: { + "constexpr_sparse_to_dense": n, + "constexpr_sparse_blockwise_shift_scale": n, + }, +} +_PALETT_EXPECTED_OPS = { + ExportBackend.CoreAI: lambda n: {"lut_to_dense": n, "sparse_to_dense": n}, + ExportBackend.CoreML: lambda n: {"constexpr_lut_to_sparse": n, "constexpr_sparse_to_dense": n}, +} + + +class TestJointQuantizationCompression: + """PTQ + PTS (post-training quantization + sparsity) across the dtype/qscheme/ + granularity matrix, on both export backends. + """ + + # int4/int8 x symmetric/asymmetric x per-tensor/per-channel. Symmetric always + # has zero_point == 0 for a signed dtype; asymmetric has a data-dependent, + # essentially-never-zero zero_point on a real trained weight -- so this + # matrix also happens to split cleanly into "accepted" / "rejected". + QUANT_VALID_CONFIGS: list[tuple[str, torch.dtype, QuantizationGranularity]] = [ + ("int8_symmetric_per_tensor", torch.int8, PerTensorGranularity()), + ("int8_symmetric_per_channel", torch.int8, PerChannelGranularity(axis=0)), + ("int4_symmetric_per_tensor", torch.int4, PerTensorGranularity()), + ("int4_symmetric_per_channel", torch.int4, PerChannelGranularity(axis=0)), + ] + QUANT_INVALID_CONFIGS: list[tuple[str, torch.dtype, QuantizationGranularity]] = [ + ("int8_asymmetric_per_tensor", torch.int8, PerTensorGranularity()), + ("int8_asymmetric_per_channel", torch.int8, PerChannelGranularity(axis=0)), + ("int4_asymmetric_per_tensor", torch.int4, PerTensorGranularity()), + ("int4_asymmetric_per_channel", torch.int4, PerChannelGranularity(axis=0)), + ] + + @staticmethod + def _build_quantizer( + model: nn.Module, + dtype: torch.dtype, + qscheme: QuantizationScheme, + granularity: QuantizationGranularity, + ) -> Quantizer: + config = QuantizerConfig( + global_config=ModuleQuantizerConfig( + op_state_spec={ + "weight": QuantizationSpec( + dtype=dtype, + qscheme=qscheme, + granularity=granularity, + _sparsity=_SPARSITY, + ) + }, + op_input_spec=None, + op_output_spec=None, + ), + execution_mode=ExecutionMode.GRAPH, + ) + return Quantizer(model, config) + + @classmethod + def _run_accepts( + cls, + backend: ExportBackend, + model: nn.Module, + input_data: torch.Tensor, + dtype: torch.dtype, + granularity: QuantizationGranularity, + expected_count: int, + ) -> None: + model.eval() + quantizer = cls._build_quantizer(model, dtype, QuantizationScheme.SYMMETRIC, granularity) + prepared_model = quantizer.prepare((input_data,)) + + with torch.no_grad(): + prepared_model_output = prepared_model(input_data) + + finalized_model = quantizer.finalize(backend=backend) + + export_utils.convert_and_verify( + finalized_model=finalized_model, + input_data=input_data, + expected_ops=_QUANT_EXPECTED_OPS[backend](expected_count), + export_backend=backend, + prepared_model_output=prepared_model_output, + ) + + @classmethod + def _run_rejects( + cls, + backend: ExportBackend, + model: nn.Module, + input_data: torch.Tensor, + dtype: torch.dtype, + granularity: QuantizationGranularity, + ) -> None: + model.eval() + quantizer = cls._build_quantizer(model, dtype, QuantizationScheme.ASYMMETRIC, granularity) + prepared_model = quantizer.prepare((input_data,)) + + with torch.no_grad(): + prepared_model(input_data) + + with pytest.raises((RuntimeError, ValueError)): + quantizer.finalize(backend=backend) + + @pytest.mark.parametrize("backend", _BACKENDS, ids=["coreai", "coreml"]) + @pytest.mark.parametrize( + "dtype,granularity", + [c[1:] for c in QUANT_VALID_CONFIGS], + ids=[c[0] for c in QUANT_VALID_CONFIGS], + ) + def test_accepts_zero_preserving_mnist( + self, backend, dtype, granularity, custom_test_mnist_model, mnist_example_input + ): + self._run_accepts( + backend, + custom_test_mnist_model, + mnist_example_input, + dtype, + granularity, + _MNIST_LAYER_COUNT, + ) + + @pytest.mark.slow + @pytest.mark.parametrize("backend", _BACKENDS, ids=["coreai", "coreml"]) + @pytest.mark.parametrize( + "dtype,granularity", + [c[1:] for c in QUANT_VALID_CONFIGS], + ids=[c[0] for c in QUANT_VALID_CONFIGS], + ) + def test_accepts_zero_preserving_resnet( + self, backend, dtype, granularity, resnet50_model, resnet_example_input + ): + self._run_accepts( + backend, resnet50_model, resnet_example_input, dtype, granularity, _RESNET_LAYER_COUNT + ) + + @pytest.mark.parametrize("backend", _BACKENDS, ids=["coreai", "coreml"]) + @pytest.mark.parametrize( + "dtype,granularity", + [c[1:] for c in QUANT_INVALID_CONFIGS], + ids=[c[0] for c in QUANT_INVALID_CONFIGS], + ) + def test_rejects_nonzero_zero_point_mnist( + self, backend, dtype, granularity, custom_test_mnist_model, mnist_example_input + ): + self._run_rejects(backend, custom_test_mnist_model, mnist_example_input, dtype, granularity) + + @pytest.mark.slow + @pytest.mark.parametrize("backend", _BACKENDS, ids=["coreai", "coreml"]) + @pytest.mark.parametrize( + "dtype,granularity", + [c[1:] for c in QUANT_INVALID_CONFIGS], + ids=[c[0] for c in QUANT_INVALID_CONFIGS], + ) + def test_rejects_nonzero_zero_point_resnet( + self, backend, dtype, granularity, resnet50_model, resnet_example_input + ): + self._run_rejects(backend, resnet50_model, resnet_example_input, dtype, granularity) + + +class TestJointPalettizationCompression: + """PTP + PTS (post-training palettization + sparsity) across the n_bits/ + cluster_dim/granularity matrix, on both export backends. + """ + + # Per-tensor, scalar (cluster_dim=1) palettization: the only combination + # joint-sparsity export supports, at a few n_bits. + PALETT_VALID_CONFIGS: list[tuple[str, dict]] = [ + ("4bit", {"n_bits": 4}), + ("6bit", {"n_bits": 6}), + ("8bit", {"n_bits": 8}), + ] + # Each violates exactly one of the two allowances (per-tensor, scalar). + PALETT_INVALID_CONFIGS: list[tuple[str, dict]] = [ + ("vector_ndim", {"n_bits": 4, "cluster_dim": 2}), + ( + "grouped_channel", + {"n_bits": 4, "granularity": PerGroupedChannelGranularity(axis=0, group_size=2)}, + ), + ] + + @staticmethod + def _build_palettizer(model: nn.Module, **spec_kwargs) -> KMeansPalettizer: + spec_kwargs.setdefault("granularity", PalettPerTensorGranularity()) + config = KMeansPalettizerConfig( + global_config=ModuleKMeansPalettizerConfig( + op_state_spec={"weight": PalettizationSpec(_sparsity=_SPARSITY, **spec_kwargs)}, + # Required whenever cluster_dim > 1 is among the configs under test. + enable_fast_kmeans_mode=False, + ) + ) + return KMeansPalettizer(model, config) + + @classmethod + def _run_accepts( + cls, + backend: ExportBackend, + model: nn.Module, + input_data: torch.Tensor, + spec_kwargs: dict, + expected_count: int, + ) -> None: + model.eval() + palettizer = cls._build_palettizer(model, **spec_kwargs) + prepared_model = palettizer.prepare((input_data,)) + + with torch.no_grad(): + prepared_model_output = prepared_model(input_data) + + finalized_model = palettizer.finalize(backend=backend) + + export_utils.convert_and_verify( + finalized_model=finalized_model, + input_data=input_data, + expected_ops=_PALETT_EXPECTED_OPS[backend](expected_count), + export_backend=backend, + prepared_model_output=prepared_model_output, + ) + + @classmethod + def _run_rejects( + cls, backend: ExportBackend, model: nn.Module, input_data: torch.Tensor, spec_kwargs: dict + ) -> None: + model.eval() + palettizer = cls._build_palettizer(model, **spec_kwargs) + prepared_model = palettizer.prepare((input_data,)) + + with torch.no_grad(): + prepared_model(input_data) + + with pytest.raises((RuntimeError, ValueError)): + palettizer.finalize(backend=backend) + + @pytest.mark.parametrize("backend", _BACKENDS, ids=["coreai", "coreml"]) + @pytest.mark.parametrize( + "spec_kwargs", + [c[1] for c in PALETT_VALID_CONFIGS], + ids=[c[0] for c in PALETT_VALID_CONFIGS], + ) + def test_accepts_scalar_per_tensor_mnist( + self, backend, spec_kwargs, custom_test_mnist_model, mnist_example_input + ): + self._run_accepts( + backend, custom_test_mnist_model, mnist_example_input, spec_kwargs, _MNIST_LAYER_COUNT + ) + + @pytest.mark.slow + @pytest.mark.parametrize("backend", _BACKENDS, ids=["coreai", "coreml"]) + @pytest.mark.parametrize( + "spec_kwargs", + [c[1] for c in PALETT_VALID_CONFIGS], + ids=[c[0] for c in PALETT_VALID_CONFIGS], + ) + def test_accepts_scalar_per_tensor_resnet( + self, backend, spec_kwargs, resnet50_model, resnet_example_input + ): + self._run_accepts( + backend, resnet50_model, resnet_example_input, spec_kwargs, _RESNET_LAYER_COUNT + ) + + @pytest.mark.parametrize("backend", _BACKENDS, ids=["coreai", "coreml"]) + @pytest.mark.parametrize( + "spec_kwargs", + [c[1] for c in PALETT_INVALID_CONFIGS], + ids=[c[0] for c in PALETT_INVALID_CONFIGS], + ) + def test_rejects_non_scalar_or_non_per_tensor_mnist( + self, backend, spec_kwargs, custom_test_mnist_model, mnist_example_input + ): + self._run_rejects(backend, custom_test_mnist_model, mnist_example_input, spec_kwargs) + + @pytest.mark.slow + @pytest.mark.parametrize("backend", _BACKENDS, ids=["coreai", "coreml"]) + @pytest.mark.parametrize( + "spec_kwargs", + [c[1] for c in PALETT_INVALID_CONFIGS], + ids=[c[0] for c in PALETT_INVALID_CONFIGS], + ) + def test_rejects_non_scalar_or_non_per_tensor_resnet( + self, backend, spec_kwargs, resnet50_model, resnet_example_input + ): + self._run_rejects(backend, resnet50_model, resnet_example_input, spec_kwargs) From 663037f1f07976fd7236a9552e364ec2585a8ab1 Mon Sep 17 00:00:00 2001 From: usimha <135899523+u-simha@users.noreply.github.com> Date: Mon, 14 Sep 2026 21:52:39 -0700 Subject: [PATCH 4/6] Update the gaurdrails to only block CoreAI path --- .../kmeans/_prepare_for_export.py | 24 ++- .../_graph/_prepare_for_export.py | 15 +- tests/export/test_joint_compression.py | 193 ++++++++++++++---- 3 files changed, 174 insertions(+), 58 deletions(-) diff --git a/src/coreai_opt/palettization/kmeans/_prepare_for_export.py b/src/coreai_opt/palettization/kmeans/_prepare_for_export.py index 81e453da..aa1d775e 100644 --- a/src/coreai_opt/palettization/kmeans/_prepare_for_export.py +++ b/src/coreai_opt/palettization/kmeans/_prepare_for_export.py @@ -166,16 +166,16 @@ def _register_mil_compression_metadata( # Determine compression type(s). coremltools' own torch-frontend converter # auto-detects sparsity from raw zeros in the traced weight value when # compression_type lists PRUNING first, then chains PALETTIZATION onto its - # constexpr_sparse_to_dense output (constexpr_lut_to_sparse) -- which has the - # same flatten-to-nonzero-only-indices contract as our own CoreAI reconstruction, - # so it needs the same per-tensor/scalar guarantee. + # constexpr_sparse_to_dense output (constexpr_lut_to_sparse). Only attempt + # that chain when it's position-independent; otherwise fall back to ordinary + # (non-joint) palettization metadata rather than rejecting the config + # outright -- unlike CoreAI, CoreML has no other combination to reject here. lut_quant = palett_info.lut_quantization if lut_quant is not None: compression_type = [CompressionType.PALETTIZATION, CompressionType.QUANTIZATION] else: compression_type = [CompressionType.PALETTIZATION] - if fake_palett_mod.sparsity is not None: - _validate_sparsity_for_export(fake_palett_mod) + if fake_palett_mod.sparsity is not None and _is_scalar_per_tensor_sparsity(fake_palett_mod): compression_type = [CompressionType.PRUNING, *compression_type] metadata = MILCompressionMetadata( @@ -242,8 +242,8 @@ def _resolve_mlir_lut_and_scale( return lut, scale, offset -def _validate_sparsity_for_export(fake_palett_mod: _FakePalettizeImplBase) -> None: - """Reject sparsity combined with anything that isn't a single, position-independent LUT. +def _is_scalar_per_tensor_sparsity(fake_palett_mod: _FakePalettizeImplBase) -> bool: + """True if masking-then-flattening the indices stays position-independent. Masking flattens indices to rank 1 before the LUT lookup, which only preserves meaning when every element shares one scalar codebook: per-tensor granularity @@ -252,6 +252,16 @@ def _validate_sparsity_for_export(fake_palett_mod: _FakePalettizeImplBase) -> No rank). A quantized LUT or per-channel scale would also apply a position-dependent mapping that flattening would scramble. """ + return ( + isinstance(fake_palett_mod.granularity, PerTensorGranularity) + and fake_palett_mod.cluster_dim == 1 + and fake_palett_mod.lut_qspec is None + and not fake_palett_mod.enable_per_channel_scale + ) + + +def _validate_sparsity_for_export(fake_palett_mod: _FakePalettizeImplBase) -> None: + """Reject sparsity combined with anything that isn't a single, position-independent LUT.""" if not isinstance(fake_palett_mod.granularity, PerTensorGranularity): raise ValueError( f"granularity={fake_palett_mod.granularity} not supported for joint sparsity." diff --git a/src/coreai_opt/quantization/_graph/_prepare_for_export.py b/src/coreai_opt/quantization/_graph/_prepare_for_export.py index 524f0b8a..ab3c11cb 100644 --- a/src/coreai_opt/quantization/_graph/_prepare_for_export.py +++ b/src/coreai_opt/quantization/_graph/_prepare_for_export.py @@ -233,13 +233,7 @@ def _register_quantization_buffers( def _validate_sparsity_for_export( fake_quant_mod: FakeQuantizeImplBase, zero_point: torch.Tensor | None ) -> None: - """Reject sparsity combined with anything ``sparse_to_dense``'s raw-0 padding can't represent. - - ``sparse_to_dense`` always pads pruned positions with a raw literal 0, so joint - sparsity is only correct when that 0 dequantizes to exactly 0.0, i.e. ``zero_point`` - is 0 (or absent). FP4 is rejected separately: its packing runs before the nonzero - values are extracted, and would misalign against the full-resolution mask. - """ + """Reject sparsity combined with anything raw-0 padding can't represent.""" if is_float4_dtype(fake_quant_mod.dtype): raise ValueError("FP4 dtype not supported for joint sparsity.") if zero_point is not None and not torch.all(zero_point == 0): @@ -502,11 +496,12 @@ def _process_mil_weight_quantization( # zeros in the registered weight value when compression_type lists PRUNING # first, then chains QUANTIZATION onto its constexpr_sparse_to_dense output # (see coremltools.converters.mil.frontend.torch.converter._construct_compression_op). - # That chain has the same raw-0-pads-pruned-positions contract as our own - # sparse_to_dense, so it needs the same zero-preserving guarantee. + # Unlike CoreAI's sparse_to_dense, that chain dequantizes the compact + # nonzero_data before scattering it into the padded dense tensor, so the + # padding is a real float 0.0, not a raw int later reinterpreted through + # zero_point -- safe for any zero_point, so there's nothing to gate here. compression_type = [CompressionType.QUANTIZATION] if fake_quant_mod.sparsity is not None: - _validate_sparsity_for_export(fake_quant_mod, zero_point) compression_type = [CompressionType.PRUNING, CompressionType.QUANTIZATION] # Create and register metadata diff --git a/tests/export/test_joint_compression.py b/tests/export/test_joint_compression.py index 13805657..aaf9da80 100644 --- a/tests/export/test_joint_compression.py +++ b/tests/export/test_joint_compression.py @@ -45,21 +45,36 @@ "constexpr_sparse_blockwise_shift_scale": n, }, } + + _PALETT_EXPECTED_OPS = { ExportBackend.CoreAI: lambda n: {"lut_to_dense": n, "sparse_to_dense": n}, ExportBackend.CoreML: lambda n: {"constexpr_lut_to_sparse": n, "constexpr_sparse_to_dense": n}, } +def _palett_ordinary_ops(n: int) -> dict[str, int]: + return {"constexpr_lut_to_dense": n} + + class TestJointQuantizationCompression: """PTQ + PTS (post-training quantization + sparsity) across the dtype/qscheme/ - granularity matrix, on both export backends. + granularity matrix. + + CoreAI's sparse_to_dense scatters raw quantized codes before dequantizing + the whole reconstructed tensor, so a nonzero zero_point would misrepresent + pruned positions -- it must reject those configs outright. CoreML's sparse + constexpr chain dequantizes the compact nonzero_data first and only then + scatters (real float 0.0 padding), which is correct for any zero_point -- + so it always takes the joint sparse chain when sparsity is set, with no + zero_point-dependent gating at all. """ # int4/int8 x symmetric/asymmetric x per-tensor/per-channel. Symmetric always # has zero_point == 0 for a signed dtype; asymmetric has a data-dependent, # essentially-never-zero zero_point on a real trained weight -- so this - # matrix also happens to split cleanly into "accepted" / "rejected". + # matrix also happens to split cleanly into "CoreAI accepts" / "CoreAI + # rejects" (CoreML accepts and joint-chains both groups identically). QUANT_VALID_CONFIGS: list[tuple[str, torch.dtype, QuantizationGranularity]] = [ ("int8_symmetric_per_tensor", torch.int8, PerTensorGranularity()), ("int8_symmetric_per_channel", torch.int8, PerChannelGranularity(axis=0)), @@ -98,17 +113,18 @@ def _build_quantizer( return Quantizer(model, config) @classmethod - def _run_accepts( + def _run( cls, backend: ExportBackend, model: nn.Module, input_data: torch.Tensor, dtype: torch.dtype, + qscheme: QuantizationScheme, granularity: QuantizationGranularity, - expected_count: int, + expected_ops: dict[str, int], ) -> None: model.eval() - quantizer = cls._build_quantizer(model, dtype, QuantizationScheme.SYMMETRIC, granularity) + quantizer = cls._build_quantizer(model, dtype, qscheme, granularity) prepared_model = quantizer.prepare((input_data,)) with torch.no_grad(): @@ -119,7 +135,7 @@ def _run_accepts( export_utils.convert_and_verify( finalized_model=finalized_model, input_data=input_data, - expected_ops=_QUANT_EXPECTED_OPS[backend](expected_count), + expected_ops=expected_ops, export_backend=backend, prepared_model_output=prepared_model_output, ) @@ -127,7 +143,6 @@ def _run_accepts( @classmethod def _run_rejects( cls, - backend: ExportBackend, model: nn.Module, input_data: torch.Tensor, dtype: torch.dtype, @@ -141,7 +156,7 @@ def _run_rejects( prepared_model(input_data) with pytest.raises((RuntimeError, ValueError)): - quantizer.finalize(backend=backend) + quantizer.finalize(backend=ExportBackend.CoreAI) @pytest.mark.parametrize("backend", _BACKENDS, ids=["coreai", "coreml"]) @pytest.mark.parametrize( @@ -152,13 +167,14 @@ def _run_rejects( def test_accepts_zero_preserving_mnist( self, backend, dtype, granularity, custom_test_mnist_model, mnist_example_input ): - self._run_accepts( + self._run( backend, custom_test_mnist_model, mnist_example_input, dtype, + QuantizationScheme.SYMMETRIC, granularity, - _MNIST_LAYER_COUNT, + _QUANT_EXPECTED_OPS[backend](_MNIST_LAYER_COUNT), ) @pytest.mark.slow @@ -171,41 +187,99 @@ def test_accepts_zero_preserving_mnist( def test_accepts_zero_preserving_resnet( self, backend, dtype, granularity, resnet50_model, resnet_example_input ): - self._run_accepts( - backend, resnet50_model, resnet_example_input, dtype, granularity, _RESNET_LAYER_COUNT + self._run( + backend, + resnet50_model, + resnet_example_input, + dtype, + QuantizationScheme.SYMMETRIC, + granularity, + _QUANT_EXPECTED_OPS[backend](_RESNET_LAYER_COUNT), ) - @pytest.mark.parametrize("backend", _BACKENDS, ids=["coreai", "coreml"]) @pytest.mark.parametrize( "dtype,granularity", [c[1:] for c in QUANT_INVALID_CONFIGS], ids=[c[0] for c in QUANT_INVALID_CONFIGS], ) - def test_rejects_nonzero_zero_point_mnist( - self, backend, dtype, granularity, custom_test_mnist_model, mnist_example_input + def test_coreai_rejects_nonzero_zero_point_mnist( + self, dtype, granularity, custom_test_mnist_model, mnist_example_input ): - self._run_rejects(backend, custom_test_mnist_model, mnist_example_input, dtype, granularity) + self._run_rejects(custom_test_mnist_model, mnist_example_input, dtype, granularity) @pytest.mark.slow - @pytest.mark.parametrize("backend", _BACKENDS, ids=["coreai", "coreml"]) @pytest.mark.parametrize( "dtype,granularity", [c[1:] for c in QUANT_INVALID_CONFIGS], ids=[c[0] for c in QUANT_INVALID_CONFIGS], ) - def test_rejects_nonzero_zero_point_resnet( - self, backend, dtype, granularity, resnet50_model, resnet_example_input + def test_coreai_rejects_nonzero_zero_point_resnet( + self, dtype, granularity, resnet50_model, resnet_example_input ): - self._run_rejects(backend, resnet50_model, resnet_example_input, dtype, granularity) + self._run_rejects(resnet50_model, resnet_example_input, dtype, granularity) + + @pytest.mark.parametrize( + "dtype,granularity", + [c[1:] for c in QUANT_INVALID_CONFIGS], + ids=[c[0] for c in QUANT_INVALID_CONFIGS], + ) + def test_coreml_accepts_nonzero_zero_point_mnist( + self, dtype, granularity, custom_test_mnist_model, mnist_example_input + ): + # Unlike CoreAI, CoreML's sparse constexpr chain dequantizes the + # compact nonzero_data before scattering it into the padded dense + # tensor, so the padding is a real float 0.0 rather than a raw int + # later reinterpreted through zero_point -- safe for any zero_point, + # so the joint chain always applies here, same op counts as symmetric. + self._run( + ExportBackend.CoreML, + custom_test_mnist_model, + mnist_example_input, + dtype, + QuantizationScheme.ASYMMETRIC, + granularity, + _QUANT_EXPECTED_OPS[ExportBackend.CoreML](_MNIST_LAYER_COUNT), + ) + + @pytest.mark.slow + @pytest.mark.parametrize( + "dtype,granularity", + [c[1:] for c in QUANT_INVALID_CONFIGS], + ids=[c[0] for c in QUANT_INVALID_CONFIGS], + ) + def test_coreml_accepts_nonzero_zero_point_resnet( + self, dtype, granularity, resnet50_model, resnet_example_input + ): + # See test_coreml_accepts_nonzero_zero_point_mnist. + self._run( + ExportBackend.CoreML, + resnet50_model, + resnet_example_input, + dtype, + QuantizationScheme.ASYMMETRIC, + granularity, + _QUANT_EXPECTED_OPS[ExportBackend.CoreML](_RESNET_LAYER_COUNT), + ) class TestJointPalettizationCompression: """PTP + PTS (post-training palettization + sparsity) across the n_bits/ - cluster_dim/granularity matrix, on both export backends. + cluster_dim/granularity matrix. + + Masking flattens indices to rank 1 before the LUT lookup, which only + preserves meaning for a single, position-independent codebook: per-tensor + granularity and scalar (cluster_dim=1) palettization. CoreAI's op chain has + no other combination to build, so it rejects unsupported configs outright. + CoreML falls back to ordinary (non-joint) compression instead of rejecting + -- but this is a genuine limitation of coremltools' own constexpr_lut_to_sparse, + not an overly-conservative gate on coreai_opt's side (unlike the quantization + zero_point case in ``TestJointQuantizationCompression``): coremltools' own + ``palettize_weights(joint_compression=True)`` falls back identically for these + same configs. """ # Per-tensor, scalar (cluster_dim=1) palettization: the only combination - # joint-sparsity export supports, at a few n_bits. + # that gets the joint sparse op chain, at a few n_bits. PALETT_VALID_CONFIGS: list[tuple[str, dict]] = [ ("4bit", {"n_bits": 4}), ("6bit", {"n_bits": 6}), @@ -233,13 +307,13 @@ def _build_palettizer(model: nn.Module, **spec_kwargs) -> KMeansPalettizer: return KMeansPalettizer(model, config) @classmethod - def _run_accepts( + def _run( cls, backend: ExportBackend, model: nn.Module, input_data: torch.Tensor, spec_kwargs: dict, - expected_count: int, + expected_ops: dict[str, int], ) -> None: model.eval() palettizer = cls._build_palettizer(model, **spec_kwargs) @@ -253,15 +327,13 @@ def _run_accepts( export_utils.convert_and_verify( finalized_model=finalized_model, input_data=input_data, - expected_ops=_PALETT_EXPECTED_OPS[backend](expected_count), + expected_ops=expected_ops, export_backend=backend, prepared_model_output=prepared_model_output, ) @classmethod - def _run_rejects( - cls, backend: ExportBackend, model: nn.Module, input_data: torch.Tensor, spec_kwargs: dict - ) -> None: + def _run_rejects(cls, model: nn.Module, input_data: torch.Tensor, spec_kwargs: dict) -> None: model.eval() palettizer = cls._build_palettizer(model, **spec_kwargs) prepared_model = palettizer.prepare((input_data,)) @@ -270,7 +342,7 @@ def _run_rejects( prepared_model(input_data) with pytest.raises((RuntimeError, ValueError)): - palettizer.finalize(backend=backend) + palettizer.finalize(backend=ExportBackend.CoreAI) @pytest.mark.parametrize("backend", _BACKENDS, ids=["coreai", "coreml"]) @pytest.mark.parametrize( @@ -281,8 +353,12 @@ def _run_rejects( def test_accepts_scalar_per_tensor_mnist( self, backend, spec_kwargs, custom_test_mnist_model, mnist_example_input ): - self._run_accepts( - backend, custom_test_mnist_model, mnist_example_input, spec_kwargs, _MNIST_LAYER_COUNT + self._run( + backend, + custom_test_mnist_model, + mnist_example_input, + spec_kwargs, + _PALETT_EXPECTED_OPS[backend](_MNIST_LAYER_COUNT), ) @pytest.mark.slow @@ -295,29 +371,64 @@ def test_accepts_scalar_per_tensor_mnist( def test_accepts_scalar_per_tensor_resnet( self, backend, spec_kwargs, resnet50_model, resnet_example_input ): - self._run_accepts( - backend, resnet50_model, resnet_example_input, spec_kwargs, _RESNET_LAYER_COUNT + self._run( + backend, + resnet50_model, + resnet_example_input, + spec_kwargs, + _PALETT_EXPECTED_OPS[backend](_RESNET_LAYER_COUNT), ) - @pytest.mark.parametrize("backend", _BACKENDS, ids=["coreai", "coreml"]) @pytest.mark.parametrize( "spec_kwargs", [c[1] for c in PALETT_INVALID_CONFIGS], ids=[c[0] for c in PALETT_INVALID_CONFIGS], ) - def test_rejects_non_scalar_or_non_per_tensor_mnist( - self, backend, spec_kwargs, custom_test_mnist_model, mnist_example_input + def test_coreai_rejects_non_scalar_or_non_per_tensor_mnist( + self, spec_kwargs, custom_test_mnist_model, mnist_example_input ): - self._run_rejects(backend, custom_test_mnist_model, mnist_example_input, spec_kwargs) + self._run_rejects(custom_test_mnist_model, mnist_example_input, spec_kwargs) @pytest.mark.slow - @pytest.mark.parametrize("backend", _BACKENDS, ids=["coreai", "coreml"]) @pytest.mark.parametrize( "spec_kwargs", [c[1] for c in PALETT_INVALID_CONFIGS], ids=[c[0] for c in PALETT_INVALID_CONFIGS], ) - def test_rejects_non_scalar_or_non_per_tensor_resnet( - self, backend, spec_kwargs, resnet50_model, resnet_example_input + def test_coreai_rejects_non_scalar_or_non_per_tensor_resnet( + self, spec_kwargs, resnet50_model, resnet_example_input + ): + self._run_rejects(resnet50_model, resnet_example_input, spec_kwargs) + + @pytest.mark.parametrize( + "spec_kwargs", + [c[1] for c in PALETT_INVALID_CONFIGS], + ids=[c[0] for c in PALETT_INVALID_CONFIGS], + ) + def test_coreml_falls_back_for_non_scalar_or_non_per_tensor_mnist( + self, spec_kwargs, custom_test_mnist_model, mnist_example_input + ): + self._run( + ExportBackend.CoreML, + custom_test_mnist_model, + mnist_example_input, + spec_kwargs, + _palett_ordinary_ops(_MNIST_LAYER_COUNT), + ) + + @pytest.mark.slow + @pytest.mark.parametrize( + "spec_kwargs", + [c[1] for c in PALETT_INVALID_CONFIGS], + ids=[c[0] for c in PALETT_INVALID_CONFIGS], + ) + def test_coreml_falls_back_for_non_scalar_or_non_per_tensor_resnet( + self, spec_kwargs, resnet50_model, resnet_example_input ): - self._run_rejects(backend, resnet50_model, resnet_example_input, spec_kwargs) + self._run( + ExportBackend.CoreML, + resnet50_model, + resnet_example_input, + spec_kwargs, + _palett_ordinary_ops(_RESNET_LAYER_COUNT), + ) From 15f9b9f57b9cdd361c256505599b9282c0f68807 Mon Sep 17 00:00:00 2001 From: usimha <135899523+u-simha@users.noreply.github.com> Date: Mon, 14 Sep 2026 22:52:35 -0700 Subject: [PATCH 5/6] Remove gating for grouped-palett --- .../kmeans/_prepare_for_export.py | 34 ++++---- tests/export/test_joint_compression.py | 82 +++++++++++++++---- 2 files changed, 82 insertions(+), 34 deletions(-) diff --git a/src/coreai_opt/palettization/kmeans/_prepare_for_export.py b/src/coreai_opt/palettization/kmeans/_prepare_for_export.py index aa1d775e..01b7f877 100644 --- a/src/coreai_opt/palettization/kmeans/_prepare_for_export.py +++ b/src/coreai_opt/palettization/kmeans/_prepare_for_export.py @@ -163,19 +163,17 @@ def _register_mil_compression_metadata( leave_parametrized=True, ) - # Determine compression type(s). coremltools' own torch-frontend converter - # auto-detects sparsity from raw zeros in the traced weight value when - # compression_type lists PRUNING first, then chains PALETTIZATION onto its - # constexpr_sparse_to_dense output (constexpr_lut_to_sparse). Only attempt - # that chain when it's position-independent; otherwise fall back to ordinary - # (non-joint) palettization metadata rather than rejecting the config - # outright -- unlike CoreAI, CoreML has no other combination to reject here. + # PRUNING must be listed first for coremltools to chain the sparse LUT op; + # only safe when the indices stay position-preserving (see + # _is_position_preserving_palettization). lut_quant = palett_info.lut_quantization if lut_quant is not None: compression_type = [CompressionType.PALETTIZATION, CompressionType.QUANTIZATION] else: compression_type = [CompressionType.PALETTIZATION] - if fake_palett_mod.sparsity is not None and _is_scalar_per_tensor_sparsity(fake_palett_mod): + if fake_palett_mod.sparsity is not None and _is_position_preserving_palettization( + fake_palett_mod + ): compression_type = [CompressionType.PRUNING, *compression_type] metadata = MILCompressionMetadata( @@ -242,19 +240,21 @@ def _resolve_mlir_lut_and_scale( return lut, scale, offset -def _is_scalar_per_tensor_sparsity(fake_palett_mod: _FakePalettizeImplBase) -> bool: +def _is_position_preserving_palettization(fake_palett_mod: _FakePalettizeImplBase) -> bool: """True if masking-then-flattening the indices stays position-independent. - Masking flattens indices to rank 1 before the LUT lookup, which only preserves - meaning when every element shares one scalar codebook: per-tensor granularity - (not per-channel/grouped, which use multiple LUTs) and cluster_dim == 1 (not - vector palettization, whose indices are already at a reduced, position-dependent - rank). A quantized LUT or per-channel scale would also apply a position-dependent - mapping that flattening would scramble. + Per-tensor and per-grouped-channel granularity both keep one index per + weight element, so coremltools' own sparse LUT op can flatten indices via + the mask and still recover per-position group context -- confirmed by + comparing against coremltools' own joint_compression flow directly. + ``cluster_dim > 1`` (vector palettization) reduces the index count below + the element count, so it can't be masked against the full-resolution mask + (coremltools itself raises an ``IndexError`` there). A quantized LUT or + per-channel scale would also apply a position-dependent mapping that + flattening would scramble. """ return ( - isinstance(fake_palett_mod.granularity, PerTensorGranularity) - and fake_palett_mod.cluster_dim == 1 + fake_palett_mod.cluster_dim == 1 and fake_palett_mod.lut_qspec is None and not fake_palett_mod.enable_per_channel_scale ) diff --git a/tests/export/test_joint_compression.py b/tests/export/test_joint_compression.py index aaf9da80..9ffeda01 100644 --- a/tests/export/test_joint_compression.py +++ b/tests/export/test_joint_compression.py @@ -266,16 +266,21 @@ class TestJointPalettizationCompression: """PTP + PTS (post-training palettization + sparsity) across the n_bits/ cluster_dim/granularity matrix. - Masking flattens indices to rank 1 before the LUT lookup, which only - preserves meaning for a single, position-independent codebook: per-tensor - granularity and scalar (cluster_dim=1) palettization. CoreAI's op chain has - no other combination to build, so it rejects unsupported configs outright. - CoreML falls back to ordinary (non-joint) compression instead of rejecting - -- but this is a genuine limitation of coremltools' own constexpr_lut_to_sparse, - not an overly-conservative gate on coreai_opt's side (unlike the quantization - zero_point case in ``TestJointQuantizationCompression``): coremltools' own - ``palettize_weights(joint_compression=True)`` falls back identically for these - same configs. + Masking flattens indices to rank 1 before the LUT lookup. CoreAI's op chain + needs a single, position-independent codebook for that to stay meaningful + -- per-tensor granularity and scalar (cluster_dim=1) -- and has no other + combination to build, so it rejects unsupported configs outright. + + CoreML is more permissive: per-tensor and per-grouped-channel granularity + both keep one index per weight element, so coremltools' own sparse LUT op + can flatten indices via the mask and still recover per-position group + context -- confirmed directly against coremltools' own + ``palettize_weights(joint_compression=True)``, so coreai_opt no longer + gates on granularity for CoreML. Vector palettization (cluster_dim>1) is + different: it reduces the index count below the element count, so + flattening against the full-resolution mask is a genuine shape mismatch -- + coremltools itself raises an ``IndexError`` there -- so CoreML still falls + back to ordinary (non-joint) compression for that one case. """ # Per-tensor, scalar (cluster_dim=1) palettization: the only combination @@ -285,7 +290,8 @@ class TestJointPalettizationCompression: ("6bit", {"n_bits": 6}), ("8bit", {"n_bits": 8}), ] - # Each violates exactly one of the two allowances (per-tensor, scalar). + # CoreAI rejects both outright. CoreML rejects only vector_ndim (see class + # docstring) -- grouped_channel gets the joint chain there. PALETT_INVALID_CONFIGS: list[tuple[str, dict]] = [ ("vector_ndim", {"n_bits": 4, "cluster_dim": 2}), ( @@ -293,6 +299,15 @@ class TestJointPalettizationCompression: {"n_bits": 4, "granularity": PerGroupedChannelGranularity(axis=0, group_size=2)}, ), ] + PALETT_COREML_FALLS_BACK_CONFIGS: list[tuple[str, dict]] = [ + ("vector_ndim", {"n_bits": 4, "cluster_dim": 2}), + ] + PALETT_COREML_ACCEPTS_CONFIGS: list[tuple[str, dict]] = [ + ( + "grouped_channel", + {"n_bits": 4, "granularity": PerGroupedChannelGranularity(axis=0, group_size=2)}, + ), + ] @staticmethod def _build_palettizer(model: nn.Module, **spec_kwargs) -> KMeansPalettizer: @@ -402,10 +417,10 @@ def test_coreai_rejects_non_scalar_or_non_per_tensor_resnet( @pytest.mark.parametrize( "spec_kwargs", - [c[1] for c in PALETT_INVALID_CONFIGS], - ids=[c[0] for c in PALETT_INVALID_CONFIGS], + [c[1] for c in PALETT_COREML_FALLS_BACK_CONFIGS], + ids=[c[0] for c in PALETT_COREML_FALLS_BACK_CONFIGS], ) - def test_coreml_falls_back_for_non_scalar_or_non_per_tensor_mnist( + def test_coreml_falls_back_for_vector_mnist( self, spec_kwargs, custom_test_mnist_model, mnist_example_input ): self._run( @@ -419,10 +434,10 @@ def test_coreml_falls_back_for_non_scalar_or_non_per_tensor_mnist( @pytest.mark.slow @pytest.mark.parametrize( "spec_kwargs", - [c[1] for c in PALETT_INVALID_CONFIGS], - ids=[c[0] for c in PALETT_INVALID_CONFIGS], + [c[1] for c in PALETT_COREML_FALLS_BACK_CONFIGS], + ids=[c[0] for c in PALETT_COREML_FALLS_BACK_CONFIGS], ) - def test_coreml_falls_back_for_non_scalar_or_non_per_tensor_resnet( + def test_coreml_falls_back_for_vector_resnet( self, spec_kwargs, resnet50_model, resnet_example_input ): self._run( @@ -432,3 +447,36 @@ def test_coreml_falls_back_for_non_scalar_or_non_per_tensor_resnet( spec_kwargs, _palett_ordinary_ops(_RESNET_LAYER_COUNT), ) + + @pytest.mark.parametrize( + "spec_kwargs", + [c[1] for c in PALETT_COREML_ACCEPTS_CONFIGS], + ids=[c[0] for c in PALETT_COREML_ACCEPTS_CONFIGS], + ) + def test_coreml_accepts_grouped_channel_mnist( + self, spec_kwargs, custom_test_mnist_model, mnist_example_input + ): + self._run( + ExportBackend.CoreML, + custom_test_mnist_model, + mnist_example_input, + spec_kwargs, + _PALETT_EXPECTED_OPS[ExportBackend.CoreML](_MNIST_LAYER_COUNT), + ) + + @pytest.mark.slow + @pytest.mark.parametrize( + "spec_kwargs", + [c[1] for c in PALETT_COREML_ACCEPTS_CONFIGS], + ids=[c[0] for c in PALETT_COREML_ACCEPTS_CONFIGS], + ) + def test_coreml_accepts_grouped_channel_resnet( + self, spec_kwargs, resnet50_model, resnet_example_input + ): + self._run( + ExportBackend.CoreML, + resnet50_model, + resnet_example_input, + spec_kwargs, + _PALETT_EXPECTED_OPS[ExportBackend.CoreML](_RESNET_LAYER_COUNT), + ) From 9e9cd532d8c1e26886a772a43ada7832b4918c3f Mon Sep 17 00:00:00 2001 From: usimha <135899523+u-simha@users.noreply.github.com> Date: Mon, 14 Sep 2026 23:13:51 -0700 Subject: [PATCH 6/6] Remove gating for lut qspec for palett --- .../kmeans/_prepare_for_export.py | 34 +++------- tests/export/test_joint_compression.py | 62 +++++++++---------- 2 files changed, 37 insertions(+), 59 deletions(-) diff --git a/src/coreai_opt/palettization/kmeans/_prepare_for_export.py b/src/coreai_opt/palettization/kmeans/_prepare_for_export.py index 01b7f877..ca9965d7 100644 --- a/src/coreai_opt/palettization/kmeans/_prepare_for_export.py +++ b/src/coreai_opt/palettization/kmeans/_prepare_for_export.py @@ -163,17 +163,19 @@ def _register_mil_compression_metadata( leave_parametrized=True, ) - # PRUNING must be listed first for coremltools to chain the sparse LUT op; - # only safe when the indices stay position-preserving (see - # _is_position_preserving_palettization). + # PRUNING must be listed first for coremltools to chain the sparse LUT op. lut_quant = palett_info.lut_quantization if lut_quant is not None: compression_type = [CompressionType.PALETTIZATION, CompressionType.QUANTIZATION] else: compression_type = [CompressionType.PALETTIZATION] - if fake_palett_mod.sparsity is not None and _is_position_preserving_palettization( - fake_palett_mod - ): + if fake_palett_mod.sparsity is not None: + # Vector palettization & per-channel scales aren't supported jointly with sparsity. + if fake_palett_mod.cluster_dim != 1 or fake_palett_mod.enable_per_channel_scale: + raise ValueError( + "cluster_dim != 1 (vector palettization) and enable_per_channel_scale " + "are not supported for joint sparsity + palettization." + ) compression_type = [CompressionType.PRUNING, *compression_type] metadata = MILCompressionMetadata( @@ -240,26 +242,6 @@ def _resolve_mlir_lut_and_scale( return lut, scale, offset -def _is_position_preserving_palettization(fake_palett_mod: _FakePalettizeImplBase) -> bool: - """True if masking-then-flattening the indices stays position-independent. - - Per-tensor and per-grouped-channel granularity both keep one index per - weight element, so coremltools' own sparse LUT op can flatten indices via - the mask and still recover per-position group context -- confirmed by - comparing against coremltools' own joint_compression flow directly. - ``cluster_dim > 1`` (vector palettization) reduces the index count below - the element count, so it can't be masked against the full-resolution mask - (coremltools itself raises an ``IndexError`` there). A quantized LUT or - per-channel scale would also apply a position-dependent mapping that - flattening would scramble. - """ - return ( - fake_palett_mod.cluster_dim == 1 - and fake_palett_mod.lut_qspec is None - and not fake_palett_mod.enable_per_channel_scale - ) - - def _validate_sparsity_for_export(fake_palett_mod: _FakePalettizeImplBase) -> None: """Reject sparsity combined with anything that isn't a single, position-independent LUT.""" if not isinstance(fake_palett_mod.granularity, PerTensorGranularity): diff --git a/tests/export/test_joint_compression.py b/tests/export/test_joint_compression.py index 9ffeda01..a9b08574 100644 --- a/tests/export/test_joint_compression.py +++ b/tests/export/test_joint_compression.py @@ -53,10 +53,6 @@ } -def _palett_ordinary_ops(n: int) -> dict[str, int]: - return {"constexpr_lut_to_dense": n} - - class TestJointQuantizationCompression: """PTQ + PTS (post-training quantization + sparsity) across the dtype/qscheme/ granularity matrix. @@ -276,11 +272,12 @@ class TestJointPalettizationCompression: can flatten indices via the mask and still recover per-position group context -- confirmed directly against coremltools' own ``palettize_weights(joint_compression=True)``, so coreai_opt no longer - gates on granularity for CoreML. Vector palettization (cluster_dim>1) is - different: it reduces the index count below the element count, so - flattening against the full-resolution mask is a genuine shape mismatch -- - coremltools itself raises an ``IndexError`` there -- so CoreML still falls - back to ordinary (non-joint) compression for that one case. + gates on granularity for CoreML. A quantized LUT (``lut_qspec``) is safe + too, since it only touches the small, fixed-size LUT array, independent of + the sparsity mask. Vector palettization (``cluster_dim>1``) and + ``enable_per_channel_scale`` are still rejected outright for CoreML: both + are genuine coremltools limitations (an index-count mismatch and a + scale/data rank mismatch, respectively), confirmed directly. """ # Per-tensor, scalar (cluster_dim=1) palettization: the only combination @@ -299,8 +296,9 @@ class TestJointPalettizationCompression: {"n_bits": 4, "granularity": PerGroupedChannelGranularity(axis=0, group_size=2)}, ), ] - PALETT_COREML_FALLS_BACK_CONFIGS: list[tuple[str, dict]] = [ + PALETT_COREML_REJECTS_CONFIGS: list[tuple[str, dict]] = [ ("vector_ndim", {"n_bits": 4, "cluster_dim": 2}), + ("per_channel_scale", {"n_bits": 4, "enable_per_channel_scale": True}), ] PALETT_COREML_ACCEPTS_CONFIGS: list[tuple[str, dict]] = [ ( @@ -348,7 +346,13 @@ def _run( ) @classmethod - def _run_rejects(cls, model: nn.Module, input_data: torch.Tensor, spec_kwargs: dict) -> None: + def _run_rejects( + cls, + backend: ExportBackend, + model: nn.Module, + input_data: torch.Tensor, + spec_kwargs: dict, + ) -> None: model.eval() palettizer = cls._build_palettizer(model, **spec_kwargs) prepared_model = palettizer.prepare((input_data,)) @@ -357,7 +361,7 @@ def _run_rejects(cls, model: nn.Module, input_data: torch.Tensor, spec_kwargs: d prepared_model(input_data) with pytest.raises((RuntimeError, ValueError)): - palettizer.finalize(backend=ExportBackend.CoreAI) + palettizer.finalize(backend=backend) @pytest.mark.parametrize("backend", _BACKENDS, ids=["coreai", "coreml"]) @pytest.mark.parametrize( @@ -402,7 +406,9 @@ def test_accepts_scalar_per_tensor_resnet( def test_coreai_rejects_non_scalar_or_non_per_tensor_mnist( self, spec_kwargs, custom_test_mnist_model, mnist_example_input ): - self._run_rejects(custom_test_mnist_model, mnist_example_input, spec_kwargs) + self._run_rejects( + ExportBackend.CoreAI, custom_test_mnist_model, mnist_example_input, spec_kwargs + ) @pytest.mark.slow @pytest.mark.parametrize( @@ -413,40 +419,30 @@ def test_coreai_rejects_non_scalar_or_non_per_tensor_mnist( def test_coreai_rejects_non_scalar_or_non_per_tensor_resnet( self, spec_kwargs, resnet50_model, resnet_example_input ): - self._run_rejects(resnet50_model, resnet_example_input, spec_kwargs) + self._run_rejects(ExportBackend.CoreAI, resnet50_model, resnet_example_input, spec_kwargs) @pytest.mark.parametrize( "spec_kwargs", - [c[1] for c in PALETT_COREML_FALLS_BACK_CONFIGS], - ids=[c[0] for c in PALETT_COREML_FALLS_BACK_CONFIGS], + [c[1] for c in PALETT_COREML_REJECTS_CONFIGS], + ids=[c[0] for c in PALETT_COREML_REJECTS_CONFIGS], ) - def test_coreml_falls_back_for_vector_mnist( + def test_coreml_rejects_vector_or_per_channel_scale_mnist( self, spec_kwargs, custom_test_mnist_model, mnist_example_input ): - self._run( - ExportBackend.CoreML, - custom_test_mnist_model, - mnist_example_input, - spec_kwargs, - _palett_ordinary_ops(_MNIST_LAYER_COUNT), + self._run_rejects( + ExportBackend.CoreML, custom_test_mnist_model, mnist_example_input, spec_kwargs ) @pytest.mark.slow @pytest.mark.parametrize( "spec_kwargs", - [c[1] for c in PALETT_COREML_FALLS_BACK_CONFIGS], - ids=[c[0] for c in PALETT_COREML_FALLS_BACK_CONFIGS], + [c[1] for c in PALETT_COREML_REJECTS_CONFIGS], + ids=[c[0] for c in PALETT_COREML_REJECTS_CONFIGS], ) - def test_coreml_falls_back_for_vector_resnet( + def test_coreml_rejects_vector_or_per_channel_scale_resnet( self, spec_kwargs, resnet50_model, resnet_example_input ): - self._run( - ExportBackend.CoreML, - resnet50_model, - resnet_example_input, - spec_kwargs, - _palett_ordinary_ops(_RESNET_LAYER_COUNT), - ) + self._run_rejects(ExportBackend.CoreML, resnet50_model, resnet_example_input, spec_kwargs) @pytest.mark.parametrize( "spec_kwargs",