diff --git a/src/coreai_opt/palettization/kmeans/_prepare_for_export.py b/src/coreai_opt/palettization/kmeans/_prepare_for_export.py index 8c0c06e2..ca9965d7 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 @@ -23,6 +24,7 @@ from coreai_opt.palettization.spec.fake_palettize import ( _FakePalettizeImplBase, ) +from coreai_opt.palettization.spec.granularity import PerTensorGranularity _DEFAULT_VECTOR_AXIS = 0 @@ -44,6 +46,34 @@ 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) + # 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 + + 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, @@ -116,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 @@ -132,12 +163,20 @@ def _register_mil_compression_metadata( leave_parametrized=True, ) - # Determine compression type(s) + # 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 + compression_type = [CompressionType.PALETTIZATION] + 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( param_name=param_name, @@ -203,11 +242,29 @@ 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.""" + 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, param_name: str, palett_info: PalettizationInfo, + fake_palett_mod: _FakePalettizeImplBase, fake_palett_idx: int, mmap_dir: str | PathLike[str] | None, ) -> None: @@ -225,6 +282,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. @@ -259,7 +320,18 @@ 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: + _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( + 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, @@ -331,10 +403,10 @@ 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_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/kmeans_fake_palettize.py b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py index 477d2e28..17c1adda 100644 --- a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py +++ b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py @@ -100,6 +100,7 @@ def __init__( rounding_precision: int = 4, op_to_optimize: Callable | None = None, training_strategy_spec: TrainingStrategySpec | None = None, + sparsity: float | None = None, ): super().__init__( n_bits=n_bits, @@ -107,6 +108,7 @@ def __init__( granularity=granularity, cluster_dim=cluster_dim, enable_per_channel_scale=enable_per_channel_scale, + sparsity=sparsity, ) self.enable_fast_kmeans_mode = enable_fast_kmeans_mode diff --git a/src/coreai_opt/palettization/kmeans/palettizer.py b/src/coreai_opt/palettization/kmeans/palettizer.py index 08a1a7ac..7c7550c2 100644 --- a/src/coreai_opt/palettization/kmeans/palettizer.py +++ b/src/coreai_opt/palettization/kmeans/palettizer.py @@ -493,6 +493,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_fake_module_kwargs()) 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 1349e28c..548fdba6 100644 --- a/src/coreai_opt/palettization/spec/fake_palettize.py +++ b/src/coreai_opt/palettization/spec/fake_palettize.py @@ -18,6 +18,7 @@ from coreai_opt.palettization.spec import ( PalettizationGranularity, ) +from coreai_opt.pruning.spec import PruneImplBase, Unstructured from coreai_opt.quantization.spec import QuantizationSpec @@ -38,6 +39,7 @@ def __init__( granularity: PalettizationGranularity, cluster_dim: int, enable_per_channel_scale: bool, + sparsity: float | None = None, **kwargs, ): super().__init__(**kwargs) @@ -46,6 +48,7 @@ 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)) # Non-persistent (kept out of new checkpoints); when set to 1 (at runtime or @@ -54,6 +57,7 @@ def __init__( "observer_enabled", torch.tensor([0], dtype=torch.uint8), persistent=False ) self._disabled = False + self.register_buffer("_sparsity_mask", None, persistent=False) self.register_buffer("indices", None) self.register_buffer("per_channel_scale", None) @@ -71,6 +75,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 + self.ensure_initialized(tensor) # Check for self._disabled again in case ensure_initialized disabled the palettizer. diff --git a/src/coreai_opt/palettization/spec/spec.py b/src/coreai_opt/palettization/spec/spec.py index 24baeb99..7148ca8c 100644 --- a/src/coreai_opt/palettization/spec/spec.py +++ b/src/coreai_opt/palettization/spec/spec.py @@ -108,6 +108,17 @@ 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._sparsity = 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 f5961328..c6d9b210 100644 --- a/src/coreai_opt/quantization/_graph/_prepare_for_export.py +++ b/src/coreai_opt/quantization/_graph/_prepare_for_export.py @@ -232,6 +232,16 @@ 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 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): + raise ValueError("nonzero zero_point not supported for joint sparsity.") + + def _process_mlir_weight_quantization( model: torch.fx.GraphModule, node: Node, @@ -269,8 +279,14 @@ def _import_coreai_custom_ops(): 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: + _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) # Drop one of the offsets so that the export @@ -286,13 +302,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: @@ -439,10 +470,22 @@ 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). + # 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: + 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/_graph/_provisional_qspec_generation.py b/src/coreai_opt/quantization/_graph/_provisional_qspec_generation.py index b1deced9..586fc720 100644 --- a/src/coreai_opt/quantization/_graph/_provisional_qspec_generation.py +++ b/src/coreai_opt/quantization/_graph/_provisional_qspec_generation.py @@ -64,6 +64,7 @@ FieldName.QPARAM_CALCULATOR_CLS: "qparam_calculator_cls", FieldName.RANGE_CALCULATOR_CLS: "range_calculator_cls", FieldName.SCALE_DTYPE: "scale_dtype", + FieldName.SPARSITY: "_sparsity", } diff --git a/src/coreai_opt/quantization/_graph/_qspec_constraints.py b/src/coreai_opt/quantization/_graph/_qspec_constraints.py index 004f875e..0425ca5a 100644 --- a/src/coreai_opt/quantization/_graph/_qspec_constraints.py +++ b/src/coreai_opt/quantization/_graph/_qspec_constraints.py @@ -320,6 +320,7 @@ def _policy_float_range_union(proposals: Sequence[FieldValue]) -> FieldValue: FieldName.QPARAM_CALCULATOR_CLS: _policy_priority_wins, FieldName.RANGE_CALCULATOR_CLS: _policy_priority_wins, FieldName.SCALE_DTYPE: _policy_priority_wins, + FieldName.SPARSITY: _policy_priority_wins, # Covering every member's values is a correctness constraint, not a # preference, so this one unions instead of deferring to priority. FieldName.FLOAT_RANGE: _policy_float_range_union, diff --git a/src/coreai_opt/quantization/_graph/_qspec_resolution.py b/src/coreai_opt/quantization/_graph/_qspec_resolution.py index e8a7d541..31ff9177 100644 --- a/src/coreai_opt/quantization/_graph/_qspec_resolution.py +++ b/src/coreai_opt/quantization/_graph/_qspec_resolution.py @@ -178,6 +178,7 @@ def _shared_spec_pointing_at(anchor: NodeSlot) -> _SharedQuantizationSpec: FieldName.QPARAM_CALCULATOR_CLS: "qparam_calculator_cls", FieldName.RANGE_CALCULATOR_CLS: "range_calculator_cls", FieldName.SCALE_DTYPE: "scale_dtype", + FieldName.SPARSITY: "_sparsity", } diff --git a/src/coreai_opt/quantization/_graph/_qspec_types.py b/src/coreai_opt/quantization/_graph/_qspec_types.py index 7c98a9d9..d509c500 100644 --- a/src/coreai_opt/quantization/_graph/_qspec_types.py +++ b/src/coreai_opt/quantization/_graph/_qspec_types.py @@ -61,6 +61,7 @@ class FieldName(enum.Enum): QPARAM_CALCULATOR_CLS = enum.auto() RANGE_CALCULATOR_CLS = enum.auto() SCALE_DTYPE = enum.auto() + SPARSITY = enum.auto() # QuantizationSpec's private, settable `_sparsity` input # Weight or activation, set by which config dict the spec came from. Not a # QuantizationSpec attribute, but an input to construct_partial alongside # them, so resolution needs it to rebuild the observer. diff --git a/src/coreai_opt/quantization/spec/factory.py b/src/coreai_opt/quantization/spec/factory.py index 408e4192..61ff004e 100644 --- a/src/coreai_opt/quantization/spec/factory.py +++ b/src/coreai_opt/quantization/spec/factory.py @@ -209,6 +209,7 @@ def create_fake_quantizer( "quant_max": spec.quant_max, "qparams_calculator": qparams_calculator, "n_bits": spec.n_bits, + "sparsity": spec._sparsity, } # Automatically detect and include any extra arguments @@ -289,6 +290,7 @@ def create_fake_quantizer_partial( "quant_min": spec.quant_min, "quant_max": spec.quant_max, "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 31b4423e..5c02cb67 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 from coreai_opt.quantization.spec.qscheme import QuantizationScheme @@ -54,6 +55,7 @@ def __init__( quant_max: int | float, qparams_calculator: QParamsCalculatorBase, n_bits: int | None = None, + sparsity: float | None = None, **kwargs, ): super().__init__() @@ -64,7 +66,9 @@ def __init__( self.quant_min = quant_min self.quant_max = quant_max self.qparams_calculator = qparams_calculator + 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: @@ -148,6 +152,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 4b3b5504..030a8bc3 100644 --- a/src/coreai_opt/quantization/spec/spec.py +++ b/src/coreai_opt/quantization/spec/spec.py @@ -374,6 +374,17 @@ 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._sparsity = sparsity + # Supported dtypes for quantization (class attribute for testing extensibility) SUPPORTED_DTYPES: ClassVar[set[torch.dtype]] = { # Signed integer types diff --git a/tests/export/test_joint_compression.py b/tests/export/test_joint_compression.py new file mode 100644 index 00000000..a9b08574 --- /dev/null +++ b/tests/export/test_joint_compression.py @@ -0,0 +1,478 @@ +# 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. + + 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 "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)), + ("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( + cls, + backend: ExportBackend, + model: nn.Module, + input_data: torch.Tensor, + dtype: torch.dtype, + qscheme: QuantizationScheme, + granularity: QuantizationGranularity, + expected_ops: dict[str, int], + ) -> None: + model.eval() + quantizer = cls._build_quantizer(model, dtype, qscheme, 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=expected_ops, + export_backend=backend, + prepared_model_output=prepared_model_output, + ) + + @classmethod + def _run_rejects( + cls, + 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=ExportBackend.CoreAI) + + @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( + backend, + custom_test_mnist_model, + mnist_example_input, + dtype, + QuantizationScheme.SYMMETRIC, + granularity, + _QUANT_EXPECTED_OPS[backend](_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( + backend, + resnet50_model, + resnet_example_input, + dtype, + QuantizationScheme.SYMMETRIC, + granularity, + _QUANT_EXPECTED_OPS[backend](_RESNET_LAYER_COUNT), + ) + + @pytest.mark.parametrize( + "dtype,granularity", + [c[1:] for c in QUANT_INVALID_CONFIGS], + ids=[c[0] for c in QUANT_INVALID_CONFIGS], + ) + def test_coreai_rejects_nonzero_zero_point_mnist( + self, dtype, granularity, custom_test_mnist_model, mnist_example_input + ): + self._run_rejects(custom_test_mnist_model, mnist_example_input, dtype, granularity) + + @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_coreai_rejects_nonzero_zero_point_resnet( + self, dtype, granularity, resnet50_model, resnet_example_input + ): + 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. + + 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. 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 + # 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}), + ("8bit", {"n_bits": 8}), + ] + # 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}), + ( + "grouped_channel", + {"n_bits": 4, "granularity": PerGroupedChannelGranularity(axis=0, group_size=2)}, + ), + ] + 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]] = [ + ( + "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( + cls, + backend: ExportBackend, + model: nn.Module, + input_data: torch.Tensor, + spec_kwargs: dict, + expected_ops: dict[str, 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=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: + 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( + backend, + custom_test_mnist_model, + mnist_example_input, + spec_kwargs, + _PALETT_EXPECTED_OPS[backend](_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( + backend, + resnet50_model, + resnet_example_input, + spec_kwargs, + _PALETT_EXPECTED_OPS[backend](_RESNET_LAYER_COUNT), + ) + + @pytest.mark.parametrize( + "spec_kwargs", + [c[1] for c in PALETT_INVALID_CONFIGS], + ids=[c[0] for c in PALETT_INVALID_CONFIGS], + ) + def test_coreai_rejects_non_scalar_or_non_per_tensor_mnist( + self, spec_kwargs, custom_test_mnist_model, mnist_example_input + ): + self._run_rejects( + ExportBackend.CoreAI, custom_test_mnist_model, mnist_example_input, spec_kwargs + ) + + @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_coreai_rejects_non_scalar_or_non_per_tensor_resnet( + self, spec_kwargs, resnet50_model, resnet_example_input + ): + self._run_rejects(ExportBackend.CoreAI, resnet50_model, resnet_example_input, spec_kwargs) + + @pytest.mark.parametrize( + "spec_kwargs", + [c[1] for c in PALETT_COREML_REJECTS_CONFIGS], + ids=[c[0] for c in PALETT_COREML_REJECTS_CONFIGS], + ) + def test_coreml_rejects_vector_or_per_channel_scale_mnist( + self, spec_kwargs, custom_test_mnist_model, mnist_example_input + ): + 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_REJECTS_CONFIGS], + ids=[c[0] for c in PALETT_COREML_REJECTS_CONFIGS], + ) + def test_coreml_rejects_vector_or_per_channel_scale_resnet( + self, spec_kwargs, resnet50_model, resnet_example_input + ): + self._run_rejects(ExportBackend.CoreML, resnet50_model, resnet_example_input, spec_kwargs) + + @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), + ) 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 + ) diff --git a/tests/quantization/test_qspec_reconcile.py b/tests/quantization/test_qspec_reconcile.py index 45df3c7c..d504399a 100644 --- a/tests/quantization/test_qspec_reconcile.py +++ b/tests/quantization/test_qspec_reconcile.py @@ -618,6 +618,7 @@ def _fields(**overrides): FieldName.QPARAM_CALCULATOR_CLS: _fv(MovingAverageQParamsCalculator), FieldName.RANGE_CALCULATOR_CLS: _fv(MinMaxRangeCalculator), FieldName.SCALE_DTYPE: _fv(None), + FieldName.SPARSITY: _fv(None), FieldName.QUANTIZATION_TARGET: _fv(CompressionTargetTensor.ACTIVATION), } # Keyword keys arrive as strings; map them onto FieldName so they