Skip to content

Add initial support for joint sparsity pipeline - #80

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

Add initial support for joint sparsity pipeline#80
u-simha wants to merge 3 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 quantization, only symmetric and integer-based quantization is supported; Only graph mode is supported
  • For palettization, only per-tensor granularity is supported; 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).
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.

1 participant