aimet_torch.experimental.spinquant

Top level APIs

aimet_torch.experimental.spinquant.apply_spinquant(model, *, enable_r1=True, enable_r2=False)

Apply SpinQuant rotation transforms to a transformer-based language model.

SpinQuant applies orthogonal Hadamard rotations to model weights to reduce quantization error. This method modifies the model in-place by:

  1. Fusing RMS normalization layers into subsequent linear layers

  2. Applying R1 Hadamard rotations to embeddings, attention, and MLP layers

  3. Applying R2 per-head Hadamard rotations to the attention value/output path

  4. Merging all transforms into the weight matrices

Supported architectures:

  • LLaMA

  • Qwen2, Qwen3

  • Phi3

  • Qwen2.5-VL (Vision-Language Model)

Parameters:

  • model (Module) – A HuggingFace transformer model (e.g., LlamaForCausalLM, Qwen2ForCausalLM). The model must have untied embed_tokens and lm_head weights.

  • enable_r1 (bool) – If True, apply the R1 (residual-stream) rotation. Defaults to True.

  • enable_r2 (bool) – If True, apply the R2 (per-head value/output) rotation. Defaults to False. Only supported on backbones with separated Q/K/V projections; not supported on fused-QKV (e.g. Phi3) or linear-attention blocks. R2 is applied to the language backbone only (not visual encoders).

Raises:

RuntimeError – If embed_tokens and lm_head weights are tied.

Example

>>> from transformers import AutoModelForCausalLM
>>> from aimet_torch.experimental.spinquant import apply_spinquant
>>> model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B-Instruct")
>>> # Untie embedding and lm_head weights if they are tied
>>> old_weight = model.lm_head.weight
>>> model.lm_head.weight = torch.nn.Parameter(
...     old_weight.data.clone().detach().to(old_weight.device),
...     requires_grad=old_weight.requires_grad,
... )
>>> apply_spinquant(model, enable_r1=True, enable_r2=True)