llmcompressor.entrypoints.model_free.spinquant

Bases: Converter

Model-free implementation of the offline SpinQuant rotations (https://arxiv.org/abs/2405.16406). Norm scales are fused into the linears which follow them, then R1 is fused into every layer which reads or writes the residual stream and R2 is fused into each attention head's value/output projections. The resulting checkpoint is numerically equivalent to the source model, but has weights which are easier to quantize. No online transforms are added, so the checkpoint loads like any dense (or subsequently quantized) model.

Construct with :meth:from_pretrained and chain before quantization, e.g.

model_free_ptq(
    MODEL_ID,
    SAVE_DIR,
    scheme="MXFP4",
    converter=[
        FP8BlockDequantizer(),
        SpinQuantConverter.from_pretrained(MODEL_ID),
    ],
)

Parameters:

• hidden_size (int) –

  size of the residual stream

• head_dim (int) –

  size of each attention head

• norms (dict[str, Tensor]) –

  module name -> norm weight, for every norm being fused

• tie_word_embeddings (bool, default: False ) –

  whether the source model ties lm_head to the embedding. If so, the saved config is untied, since the rotated lm_head and embedding differ

• emit_lm_head (bool, default: False ) –

  whether to create lm_head.weight from the embedding, for tied checkpoints which do not save lm_head.weight

• rotations (tuple[Literal['R1', 'R2'], ...], default: ('R1', 'R2') ) –

  offline rotations to apply. from_pretrained defaults an MLA model (one with kv_lora_rank) to R1 only: R2 needs an uncompressed num_heads * head_dim path between v_proj and o_proj, which MLA replaces with a compressed latent

• transform_type (Literal['hadamard', 'random-hadamard'], default: 'hadamard' ) –

  "hadamard" supports power-of-two sizes, "random-hadamard" also supports sizes with a known hadamard divisor

• transform_block_size (int | None, default: None ) –

  R1 block size. Defaults to hidden_size. The rotation is block-diagonal when smaller than hidden_size

• seed (int, default: 0 ) –

  seed for "random-hadamard"

• precision (dtype, default: float64 ) –

  dtype used to fuse rotations into weights

• mapping (SpinQuantConverterMapping | None, default: None ) –

  names of the layers to rotate

• ignore (list[str], default: () ) –

  regexes of tensor names which are left untouched, for modules outside of the rotated residual stream such as a vision tower

Methods:

• from_pretrained –

  Build the converter from a checkpoint's config.json and norm weights. Norm

146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210

def __init__(
    self,
    hidden_size: int,
    head_dim: int,
    norms: dict[str, torch.Tensor],
    tie_word_embeddings: bool = False,
    emit_lm_head: bool = False,
    rotations: tuple[Literal["R1", "R2"], ...] = ("R1", "R2"),
    transform_type: Literal["hadamard", "random-hadamard"] = "hadamard",
    transform_block_size: int | None = None,
    seed: int = 0,
    precision: torch.dtype = torch.float64,
    mapping: SpinQuantConverterMapping | None = None,
    ignore: list[str] = (),
    num_attention_heads: int | None = None,
    num_key_value_heads: int | None = None,
):
    rotations = tuple(r.upper() for r in rotations)
    if not rotations or not set(rotations) <= {"R1", "R2"}:
        raise ValueError(
            f"SpinQuantConverter supports R1 and R2 only, got {rotations}"
        )

    self.hidden_size = hidden_size
    self.head_dim = head_dim
    self.norms = norms
    self.tie_word_embeddings = tie_word_embeddings
    self.emit_lm_head = emit_lm_head
    self.rotations = rotations
    self.precision = precision
    self.mapping = mapping or SpinQuantConverterMapping()

    block_size = transform_block_size or hidden_size
    if hidden_size % block_size != 0:
        raise ValueError(
            f"transform_block_size {block_size} must divide "
            f"hidden_size {hidden_size}"
        )
    generator = torch.Generator().manual_seed(seed)
    self.r1 = _hadamard(block_size, transform_type, precision, generator)
    # R2's head-sized matrix is only needed, and only constructible for every
    # head_dim, when R2 is actually applied.
    self.r2 = (
        _hadamard(head_dim, transform_type, precision, generator)
        if "R2" in rotations
        else None
    )
    self.r2_v_dim = (
        None if num_key_value_heads is None else num_key_value_heads * head_dim
    )
    self.r2_o_dim = (
        None if num_attention_heads is None else num_attention_heads * head_dim
    )

    m = self.mapping
    self._embedding = re.compile(m.embedding)
    self._lm_head = re.compile(m.lm_head)
    self._attn_in = [re.compile(p) for p in m.attn_in]
    self._attn_out = [re.compile(p) for p in m.attn_out]
    self._mlp_in = [re.compile(p) for p in m.mlp_in]
    self._mlp_out = [re.compile(p) for p in m.mlp_out]
    self._attn_v = re.compile(m.attn_v)
    self._attn_o = re.compile(m.attn_o)
    self._residual = [re.compile(p) for p in m.residual]
    self._ignore = [re.compile(p) for p in ignore]

from_pretrained(
    model_name_or_path: str | PathLike, **kwargs
) -> SpinQuantConverter

212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278

@classmethod
def from_pretrained(
    cls, model_name_or_path: str | os.PathLike, **kwargs
) -> "SpinQuantConverter":
    """
    Build the converter from a checkpoint's config.json and norm weights. Norm
    weights are small, so they are loaded up front rather than declared as
    dependencies of every linear which reads them.

    :param model_name_or_path: HuggingFace stub or local checkpoint path
    :param kwargs: additional arguments passed to the constructor
    """
    model_files = get_checkpoint_files(model_name_or_path)
    weight_map = get_weight_map(model_files)
    config = load_model_config(model_files)
    text_config = config.get("text_config", config)

    hidden_size = text_config["hidden_size"]
    head_dim = text_config.get("head_dim") or (
        hidden_size // text_config["num_attention_heads"]
    )
    tie_word_embeddings = config.get(
        "tie_word_embeddings", text_config.get("tie_word_embeddings", False)
    )

    mapping = kwargs.get("mapping") or SpinQuantConverterMapping()
    ignore = [re.compile(p) for p in kwargs.get("ignore", ())]
    names = [name for name in weight_map if not _search_any(ignore, name)]
    lm_head = re.compile(mapping.lm_head)
    has_lm_head = any(
        lm_head.search(name.removesuffix(".weight"))
        for name in names
        if name.endswith(".weight")
    )

    norm_names = _find_norm_names(names, mapping)
    if not norm_names:
        raise ValueError("No norm weights found to fuse into linear layers")

    names_by_file: dict[str, list[str]] = defaultdict(list)
    for name in norm_names:
        names_by_file[model_files[weight_map[name]]].append(name)

    norms = {}
    for path, names in names_by_file.items():
        with safe_open(path, framework="pt") as file:
            for name in names:
                norms[name.removesuffix(".weight")] = file.get_tensor(name)

    if "rotations" not in kwargs and _is_mla(text_config):
        kwargs["rotations"] = ("R1",)
        logger.warning(
            "MLA model detected (kv_lora_rank is set). SpinQuant R2 needs an "
            "uncompressed head dimension between v_proj and o_proj, so "
            "rotations default to R1 only. Pass rotations explicitly to override"
        )
    kwargs.setdefault("num_attention_heads", text_config.get("num_attention_heads"))
    kwargs.setdefault("num_key_value_heads", text_config.get("num_key_value_heads"))

    return cls(
        hidden_size=hidden_size,
        head_dim=head_dim,
        norms=norms,
        tie_word_embeddings=tie_word_embeddings,
        emit_lm_head=tie_word_embeddings and not has_lm_head,
        **kwargs,
    )

Bases: BaseModel

Name-based description of where SpinQuant rotations are fused in a checkpoint. Patterns are regexes searched against module names (tensor names without the trailing .weight / .bias). Defaults target pre-norm decoders with Llama-style names, including MoE routers and experts.

Parameters:

• embedding –

  embedding whose rows are rotated by R1

• attn_in –

  attention layers reading the residual stream

• attn_out –

  attention layers writing the residual stream

• attn_v –

  layer whose output heads are rotated by R2

• attn_o –

  layer whose input heads are rotated by R2

• mlp_in –

  mlp layers reading the residual stream

• mlp_out –

  mlp layers writing the residual stream

• lm_head –

  language model head

• residual –

  bare parameters, without a .weight suffix, whose last dimension is one or more residual streams. Each stream is rotated by R1, i.e. a last dimension of k * hidden_size is rotated by a block-diagonal R1. Defaults include DeepSeek-V4 mHC mixers (hc_attn_fn, hc_ffn_fn, hc_head_fn)

• attn_norm –

  norm fused into attn_in, relative to the decoder layer

• mlp_norm –

  norm fused into mlp_in, relative to the decoder layer

• final_norm –

  norm fused into lm_head, relative to the decoder prefix