diff --git a/src/coreai_opt/palettization/kmeans/_prepare_for_export.py b/src/coreai_opt/palettization/kmeans/_prepare_for_export.py index 8c0c06e..a6172ff 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 @@ -44,6 +45,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, @@ -208,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: @@ -225,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. @@ -259,7 +293,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: + # 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( + 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, @@ -334,7 +379,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/kmeans_fake_palettize.py b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py index 566c140..51c9bee 100644 --- a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py +++ b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py @@ -82,6 +82,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, @@ -89,6 +90,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 b85d7c6..20cc963 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 34dff09..47c66bd 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 24baeb9..9574915 100644 --- a/src/coreai_opt/palettization/spec/spec.py +++ b/src/coreai_opt/palettization/spec/spec.py @@ -108,6 +108,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: + 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): + 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 477b03f..9da0e66 100644 --- a/src/coreai_opt/quantization/_graph/_prepare_for_export.py +++ b/src/coreai_opt/quantization/_graph/_prepare_for_export.py @@ -267,8 +267,13 @@ 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: + 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 @@ -284,13 +289,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/_graph/_provisional_qspec_generation.py b/src/coreai_opt/quantization/_graph/_provisional_qspec_generation.py index b1deced..586fc72 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 004f875..0425ca5 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 e8a7d54..31ff917 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 7c98a9d..d509c50 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 408e419..61ff004 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 9ff72aa..27eb87f 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: @@ -146,6 +150,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 c39f3a1..0495c2f 100644 --- a/src/coreai_opt/quantization/spec/spec.py +++ b/src/coreai_opt/quantization/spec/spec.py @@ -373,6 +373,16 @@ 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: + 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 @@ -556,6 +566,22 @@ 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_sparsity.py b/tests/export/test_joint_sparsity.py new file mode 100644 index 0000000..5957d59 --- /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 45df3c7..d504399 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