sparsification#

High-level API to sparsify the weights and the attention of a model.

Functions

calibrate

Calibrates sparse attention thresholds based on target sparsity.

export

Export a sparse model to a regular model by folding the weight sparsity masks into the weights.

sparsify

Sparsify the weights and/or the attention of a model in-place.
calibrate(model, config, forward_loop=None)#

Calibrates sparse attention thresholds based on target sparsity.

Parameters:

  • model (Module) – Model with sparse attention modules

  • config (dict[str, Any] | SparseAttentionConfig) – Sparse attention configuration with calibration settings

  • forward_loop (Callable[[Module], None] | None) – Optional callable that forwards calibration data through the model. If provided, uses this for calibration data. If None, will auto-generate RULER dataset for calibration.

Returns:

The calibrated model with optimized sparse attention thresholds.

Return type:

Module

export(model)#

Export a sparse model to a regular model by folding the weight sparsity masks into the weights.

Export after fine-tuning: the masks are no longer enforced afterwards. Sparse attention needs no export and stays applied.

Parameters:

model (Module)

Return type:

Module

sparsify(model, mode, config=None, *, forward_loop=None)#

Sparsify the weights and/or the attention of a model in-place.

Example:

model = mts.sparsify(model, "sparsegpt", config={"data_loader": data_loader})
model = mts.sparsify(model, mts.SKIP_SOFTMAX_CALIB)  # calibrates on RULER data
Parameters:

  • model (Module) – The model to sparsify.

  • mode (ModeDescriptor | str | list[ModeDescriptor | str] | list[tuple[str, dict[str, Any]]] | dict[str, Any] | SparseAttentionConfig) –

    A mode, a list of modes, or a list of (mode, mode_config) tuples. Modes are applied and then searched or calibrated one at a time, in order:

    • "sparse_magnitude": 2:4 weight sparsity by weight magnitude.

    • "sparsegpt": 2:4 weight sparsity with SparseGPT weight updates.

    • "sparse_attention": sparse attention configured by a SparseAttentionConfig such as mts.SKIP_SOFTMAX_CALIB, which may also be passed directly as mode.

  • config (dict[str, Any] | None) – Search options such as forward_loop, data_loader, collect_func and verbose, plus algorithm-specific ones such as SparseGPT’s device. Each mode receives the options its algorithm accepts.

  • forward_loop (Callable[[Module], None] | None) – Shorthand for config["forward_loop"]. Sparse attention calibration accepts a forward_loop or data_loader only with a decode target sparsity of 0; without either, it calibrates on generated RULER data.

Returns:

The sparsified model.

Return type:

Module