Skip to content

Add initial support for joint sparsity pipeline - #80

Open
u-simha wants to merge 10 commits into
apple:mainfrom
u-simha:u/usimha/add-quant-plus-sparsity-support
Open

u-simha wants to merge 10 commits into
apple:mainfrom
u-simha:u/usimha/add-quant-plus-sparsity-support

Conversation

@u-simha

@u-simha u-simha commented Aug 28, 2026

Copy link
Copy Markdown
Contributor

What:

  • Add support for testing out sparsity with quantization & palettization
  • This is mostly for verifying export support for joint sparsity

How:

  • Add a private _sparsity field to the QuantizationSpec and PalettizationSpec to sparsity the weights before the fake quantization & fake palettization is applied
  • Update prepare_for_export for both quantization & palettization to insert the sparse_to_dense ops in the correct order
  • Add some basic end-to-end testing to validate the ops are inserted and the SNR of the prepared model and the exported model match

Constraints:

  • Supported for post-training quantization & post-training palettization (we use the mask computed in the prepare() call)
  • For CoreAI backend,
    • For quantization, only zero point-based and integer-based quantization is supported; Only graph mode is supported
    • For palettization, only per-tensor granularity is supported; Vector palettization, quantized LUT and per-channel scales are not supported

Example quant + sparsity config:

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,
        )

Example palett + sparsity config:

config = KMeansPalettizerConfig(
            global_config=ModuleKMeansPalettizerConfig(
                op_state_spec={"weight": PalettizationSpec(n_bits=8, _sparsity=0.5)}
            )
        )

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.
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.
Resolves conflicts from main's graph-mode rewrite (qspec reconciliation
pipeline, ensure_initialized-based palettization clustering, factory/import
renames) and threads the sparsity plumbing through the new code paths:

- _KMeansFakePalettize.__init__ now takes sparsity explicitly and forwards it
  to the base class, since upstream removed the old **kwargs passthrough.
- The graph-mode qspec reconciliation pipeline (_qspec_types.py,
  _provisional_qspec_generation.py, _qspec_constraints.py,
  _qspec_resolution.py) rebuilds a fresh QuantizationSpec from reconciled
  fields; added a SPARSITY field so the private _sparsity value round-trips
  through that rebuild instead of silently reverting to None.

Verified: tests/export/test_joint_sparsity.py (4/4), tests/quantization/ +
tests/palettization/ (2059/2059, excluding slow), tests/export/ full suite
(all passing, excluding slow).
…rse 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.
@u-simha
u-simha marked this pull request as ready for review September 11, 2026 22:51
…-support' into u/usimha/add-quant-plus-sparsity-support
FieldName.QPARAM_CALCULATOR_CLS: "qparam_calculator_cls",
FieldName.RANGE_CALCULATOR_CLS: "range_calculator_cls",
FieldName.SCALE_DTYPE: "scale_dtype",
FieldName.SPARSITY: "_sparsity",

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@anotheranshu - I realized adding a new spec value is pretty cumbersome; is there a way we can have this as **args or something? There are 4 files I needed to update, which is not ideal.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants