SpinQuant

Note

This feature is currently experimental. The API may change in the future.

Context

SpinQuant (arXiv:2405.16406) is a post-training quantization technique that reduces activation outliers by inserting orthogonal Hadamard rotations at key points in the model. Because the rotations are absorbed into adjacent weight matrices, the final model architecture is unchanged.

AIMET implements R1 and R2 rotations (fixed Hadamard, no optimization).

  • R1 fuses RMSNorm scale weights into downstream linear layers, then applies a Hadamard rotation across the residual stream to reduce outliers in Q/K/V/O and gate/up/down projections. Enabled by default.

  • R2 applies a per-head Hadamard rotation across the attention head dimension, absorbed into V and O projections to reduce outliers in the attention output. Disabled by default; enable via enable_r2=True. Not supported on architectures with fused QKV projections (e.g. Phi3).

The ONNX apply_spinquant API exposes these as boolean flags enable_r1 and enable_r2.

Supported architectures

Model family

PyTorch

ONNX

LlamaForCausalLM

✓

✓

Qwen2ForCausalLM, Qwen3ForCausalLM

✓

✓

MistralForCausalLM

✓

✓

Phi3ForCausalLM

✓

✓

Qwen2.5-VL, Qwen3-VL

✓

✓

Note

Support for additional model families is added continuously as new architectures are validated. See the release notes for the latest additions.

Workflow

Prerequisites

To use SpinQuant, you need:

  • A pre-trained model loaded from HuggingFace.

  • ONNX only: the model must be exported to ONNX — Step 2 handles this.

Note

For a complete working example, see Examples/torch/quantize.py or Examples/onnx/quantize.py (run with --recipe pcq_spinquant).

Procedure

Step 1: Load model

Load the HuggingFace model and wrap it with ONNXExportableModuleWithCache to enable JIT tracing with a static graph.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from GenAILab.qai_hub_lm.models.utils.exportable import ONNXExportableModuleWithCache

SEQUENCE_LENGTH = 2048
CONTEXT_LENGTH = 4096

model_id = "meta-llama/Llama-3.2-1B-Instruct"
hf_model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.float32)
tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True, trust_remote_code=True)

# Wrap model to satisfy static graph constraints for JIT trace
traceable_model = ONNXExportableModuleWithCache(hf_model)

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from GenAILab.qai_hub_lm.models.utils.exportable import ONNXExportableModuleWithCache

SEQUENCE_LENGTH = 2048
CONTEXT_LENGTH = 4096

model_id = "meta-llama/Llama-3.2-1B-Instruct"
hf_model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.float32)
tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True, trust_remote_code=True)

# Wrap model to satisfy static graph constraints for JIT trace
traceable_model = ONNXExportableModuleWithCache(hf_model)

Step 2: Create QuantizationSimModel

Create a QuantizationSimModel with the desired quantization configuration.

ONNX only: this step instead exports the model to ONNX and analyzes its decoder-stack structure with analyze_llm_topology. The sim is created in Step 3, after the rotation, because ONNX SpinQuant rewrites the float graph and R3 inserts new ops that only a sim built on the rotated graph can wrap in quantizers. A topology describes the model, not the sim, so derive it once here; Step 3 hands it to SpinQuant, which uses it to place every rotation.

from aimet_torch.common.defs import QuantScheme
from aimet_torch import QuantizationSimModel
from GenAILab.qai_hub_lm.models.generator import Generator

assembled_dummy_inputs = tuple(
    Generator.prepare_inputs(
        model=traceable_model,
        input_ids=torch.zeros((1, SEQUENCE_LENGTH), dtype=torch.int),
        attention_mask=torch.ones((1, SEQUENCE_LENGTH), dtype=torch.int),
        past_key_values=[],
        context_length=CONTEXT_LENGTH,
        sequence_length=SEQUENCE_LENGTH,
    ).values()
)

quantsim = QuantizationSimModel(
    model=traceable_model,
    quant_scheme=QuantScheme.post_training_tf,
    dummy_input=assembled_dummy_inputs,
    default_output_bw=16,
    default_param_bw=8,
    in_place=True,
    config_file="htp_v73",
)

generator = Generator(quantsim.model, tokenizer, SEQUENCE_LENGTH, CONTEXT_LENGTH)

import os
import tempfile
import onnx
from aimet_onnx.experimental.llm_topology import analyze_llm_topology
from GenAILab.qai_hub_lm.models.base import LLM
from GenAILab.qai_hub_lm.models.utils.layer_cache import build_layer_cache_descriptors
from GenAILab.qai_hub_lm.models.generator import Generator

assembled_dummy_inputs = tuple(
    Generator.prepare_inputs(
        model=traceable_model,
        input_ids=torch.zeros((1, SEQUENCE_LENGTH), dtype=torch.int),
        attention_mask=torch.ones((1, SEQUENCE_LENGTH), dtype=torch.int),
        past_key_values=[],
        context_length=CONTEXT_LENGTH,
        sequence_length=SEQUENCE_LENGTH,
    ).values()
)

layer_cache_descs = build_layer_cache_descriptors(hf_model.config)
with tempfile.TemporaryDirectory() as tmpdir:
    torch.onnx.export(
        traceable_model,
        assembled_dummy_inputs,
        os.path.join(tmpdir, "model.onnx"),
        input_names=LLM.get_backbone_input_names(layer_cache_descs),
        output_names=LLM.get_backbone_output_names(layer_cache_descs),
        opset_version=17,
        dynamo=False,
    )
    onnx_model = onnx.load(os.path.join(tmpdir, "model.onnx"))

# Analyze the decoder-stack structure. The topology describes the model, so it is
# derived once here and then handed to any technique that needs to know where the
# blocks and their projections are (SpinQuant, below).
topology = analyze_llm_topology(onnx_model)

Step 3: Apply SpinQuant

Apply SpinQuant to the model. This fuses RMSNorm scale weights into downstream linear layers and applies the R1 Hadamard rotation to all weight matrices in-place.

Important

apply_spinquant must be called before compute_encodings. The rotation modifies float weight initializers; compute_encodings must run afterward to calibrate quantizer scales on the rotated weights. For ONNX it must also run before the sim is created, so the ONNX tab below rotates the float model and then builds the sim on the rotated graph.

Note

ONNX only: SpinQuant is told where each rotation goes, rather than finding out itself. Pass the LlmTopology analyzed in Step 2 as topology — locating model structure belongs to AIMET’s topology analysis, and SpinQuant only rotates what it reports. The argument is optional today: omitting it analyzes the topology internally and warns. Passing it is the recommended usage, and may become required in a future release.

from aimet_torch.experimental.spinquant import apply_spinquant

# apply_spinquant modifies the model in-place. Must be called BEFORE compute_encodings.
apply_spinquant(model=quantsim.model)

from aimet_onnx.experimental.spinquant import apply_spinquant

# apply_spinquant rotates onnx_model in-place. Must be called on the float model,
# BEFORE the sim is created: R3 inserts new ops that only a sim built afterward can
# wrap in quantizers. ``topology`` tells it where each rotation goes.
apply_spinquant(onnx_model, topology=topology)

from aimet_onnx.quantsim import QuantizationSimModel
from GenAILab.qai_hub_lm.backends.onnx.torch_onnx_interface import TorchONNXInterface
from GenAILab.qai_hub_lm.backends.onnx.quantsim_utils import (
    _set_tensors_to_output_n_bit_symmmetric,
    _tie_quantizers_for_kv_cache,
    _set_lm_head_precision,
)
from GenAILab.bench.precision import WeightPrecision
from aimet_onnx.common.defs import int8

# Built on the rotated graph, so compute_encodings calibrates on rotated weights.
quantsim = QuantizationSimModel(
    model=onnx_model,
    quant_scheme="min_max",
    default_activation_bw=16,
    default_param_bw=8,
    config_file="htp_v73",
    providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
)
_set_tensors_to_output_n_bit_symmmetric(quantsim, kv_bits=8)
_set_lm_head_precision(quantsim, WeightPrecision(qtype=int8, granularity="PCQ"))
_tie_quantizers_for_kv_cache(quantsim)

quantsim_with_torch_interface = TorchONNXInterface(quantsim, hf_model.config)
generator = Generator(quantsim_with_torch_interface, tokenizer, SEQUENCE_LENGTH, CONTEXT_LENGTH)

Step 4: Compute activation encodings

Calibrate activation quantizers by running the model through a representative dataset.

import itertools
from tqdm import tqdm
from GenAILab.bench.datasets import Wikitext

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
train_dataset = Wikitext.load_encoded_dataset(tokenizer, CONTEXT_LENGTH, "train")


def calibration_callback(model):
    for batch in tqdm(
        itertools.islice(train_dataset, 20), total=20, desc="Calibrating"
    ):
        generator(input_ids=batch["input_ids"].to(device=device))


quantsim.compute_encodings(calibration_callback)

from tqdm import tqdm
from GenAILab.bench.datasets import Wikitext
from GenAILab.bench.onnx.quant_recipes import _prefill_inputs

train_dataset = Wikitext.load_encoded_dataset(tokenizer, CONTEXT_LENGTH, "train")
calib_inputs = _prefill_inputs(quantsim, generator, train_dataset, num_batches=20)


def _forward(session, _):
    for batch in tqdm(calib_inputs, total=len(calib_inputs), desc="Calibrating"):
        session.run(None, batch)


quantsim.compute_encodings(_forward, tuple())

After completing these steps, export the quantized model:

  • PyTorch: quantsim.export(...)

  • ONNX: quantsim.export(...)

API

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)

Top level APIs

aimet_onnx.experimental.spinquant.apply_spinquant(model, visual_model=None, embedding=None, *, enable_r1=True, enable_r2=False, enable_r3=False, topology=None)[source]

Apply SpinQuant rotation transforms to an ONNX transformer model.

SpinQuant applies orthogonal Hadamard rotations to model weights to reduce quantization error. This function modifies the ONNX model(s) in-place by:

  1. Reading the backbone’s decoder-stack structure (block boundaries, role map, hidden size) off topology, and analyzing the optional visual encoder (PatchMerger output projection).

  2. Validating every selected rotation pass against the analysis.

  3. Applying every selected pass in order (R1 before R2 before R3).

  4. Serializing the result back onto the caller’s ModelProtos.

The rotations are performed on an onnx_ir copy of the graph and written back only once every pass has succeeded, so a failure part-way through leaves the caller’s model untouched rather than half-rotated. (The embedding tensor is the exception: it is a torch.Tensor rotated in place.)

Must be called on the float ONNX model BEFORE creating a QuantizationSimModel. The rotation modifies float weight initializers (R1 / R2) and may insert new nodes (R3); build the sim on the rotated graph and run compute_encodings afterward so quantizer scales are calibrated on the rotated weights.

Supported architectures:

  • LLaMA, Qwen2, Qwen3, Phi3 (backbone only)

  • Qwen2.5-VL, Qwen3-VL (backbone + visual)

Parameters:

  • model (ModelProto) – backbone.onnx ModelProto. Mutated in-place.

  • visual_model (Optional[ModelProto]) – Optional visual.onnx ModelProto (VLM only). Mutated in-place.

  • embedding (Optional[Tensor]) – Optional torch.Tensor of shape [vocab, hidden] loaded from embedding.pth (VLM only). Rotated in-place with R_L.

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

  • enable_r2 (bool) – If True, apply the R2 (per-head) rotation. Defaults to False. Not supported on architectures with fused QKV projections (e.g. Phi3).

  • enable_r3 (bool) – If True, apply the R3 online Hadamard rotation on Q and K paths into each block’s QK^T MatMul. Defaults to False. Inserts new MatMul nodes into the ONNX graph (does not mutate existing weights). MHA only — not supported on fused QKV or per-head split exports. The K-side rotation is placed upstream of the past-key Concat so K values entering the KV cache are already rotated; the model’s present_key output then carries rotated K (cache convention is self-consistent across steps).

  • topology (Optional[LlmTopology]) – Decoder-stack topology of model, from analyze_llm_topology(). Every rotation is placed by it — which linears read from and write to the residual stream (R1), which are the V/O projections (R2), which are the Q/K edges into QKᵀ (R3) — because discovering model structure belongs to llm_topology, not to SpinQuant. Analyze the same model this call rotates. Passing it explicitly is also the only way to override the analysis: this function would otherwise analyze with active_norms_per_block=2 and the default role patterns, which no argument here can change. Optional today: when omitted the topology is analyzed internally from model and a warning is raised. Passing it explicitly is recommended, and may become required in a future release.

Raises:

ValueError – If no rotation is enabled, if block detection or role classification fails, if topology does not describe model, or if any expected weight is missing / has the wrong shape.

Return type:

None

Example (LLM):

topology = analyze_llm_topology(model)
apply_spinquant(model, topology=topology)
sim = QuantizationSimModel(model)   # built on the rotated graph
sim.compute_encodings(calibration_data)

Example (VLM):

embedding = torch.load("embedding.pth")   # torch.Tensor [vocab, hidden]
topology = analyze_llm_topology(backbone_model)
apply_spinquant(
    backbone_model,
    visual_model=visual_model,
    embedding=embedding,
    topology=topology,
)
torch.save(embedding, "embedding.pth")    # overwrite with rotated weights
backbone_sim = QuantizationSimModel(backbone_model)
visual_sim = QuantizationSimModel(visual_model)
backbone_sim.compute_encodings(backbone_calibration_data)
visual_sim.compute_encodings(visual_calibration_data)

Model topology

SpinQuant places every rotation using an LlmTopology, which describes the structure of the decoder stack: which linears read from and write to the residual stream (R1), which are the V and O projections (R2), and which Q/K edges feed each QKᵀ MatMul (R3). Analyze the model, then pass the result as topology:

topology = analyze_llm_topology(onnx_model)
apply_spinquant(onnx_model, topology=topology)
sim = QuantizationSimModel(onnx_model, ...)   # built on the rotated graph
aimet_onnx.experimental.llm_topology.analyze_llm_topology(model, active_norms_per_block=2, expected_num_blocks=None, role_patterns=None, ir_model=None)[source]

Analyze model end-to-end and return a name-based LlmTopology.

Runs the whole pipeline on a private onnx_ir copy of model: strip quantizers, fuse RMSNorms, detect active norms and block boundaries, build the per-block topology, and infer hidden_size / head_dim.

Parameters:

  • model (ModelProto) – ONNX ModelProto to analyze. Not mutated.

  • active_norms_per_block (int) – Active norms per decoder block (see get_decoder_block_boundaries()). Defaults to 2.

  • expected_num_blocks (Optional[int]) – If given, validated against the detected count.

  • role_patterns (Optional[Dict[LinearRole, Pattern]]) – Optional module-name → role override (see classify_linear_role()).

  • ir_model (Optional[Model]) – Pre-built analysis IR for model — as returned by build_analysis_ir(), i.e. quantizer-stripped and RMSNorm-fused. Built here when None. Pass one only to avoid a second from_proto of a large model when the caller already holds it; a faithful (unfused) IR will not detect norms and must not be passed.

Return type:

LlmTopology

Returns:

LlmTopology with block/backbone roles, active_norms, hidden_size and head_dim populated. head_dim is None when the export exposes no past_value graph input to derive it from.

class aimet_onnx.experimental.llm_topology.LlmTopology(embed_tokens=<factory>, lm_head=<factory>, blocks=<factory>, past_key_input_names=<factory>, past_key_output_names=<factory>, past_value_input_names=<factory>, past_value_output_names=<factory>, active_norms=None, hidden_size=None, head_dim=None)[source]

Topology of an ONNX decoder-stack model: blocks + backbone-level roles + dims.

Parameters:

  • embed_tokens (List[str]) – Names of the token-embedding ``Gather``(s) that produce the initial residual-stream activations.

  • lm_head (List[str]) – Names of the vocabulary-projection linear(s) downstream of the final norm.

  • blocks (List[BlockTopology]) – Per-decoder-block topology in topological order.

  • past_key_input_names (List[str]) – Raw past_key_* graph inputs in declaration order, collected tolerantly (empty for prefill-only exports without a KV-cache). Pairing these to blocks and validating that their count matches the block count are the consumer’s responsibility (e.g. R3) — R1-only and prefill-only flows do not require KV-cache inputs.

  • past_key_output_names (List[str]) – Key-cache graph outputs in declaration order.

  • past_value_input_names (List[str]) – Value-cache graph inputs in declaration order.

  • past_value_output_names (List[str]) – Value-cache graph outputs in declaration order.

  • active_norms (Optional[List[ActiveNorm]]) – Active norms in topological order used to build the topology.

  • hidden_size (Optional[int]) – Residual-stream hidden dimension (None if not inferred; analyze_llm_topology() fills it).

  • head_dim (Optional[int]) – Per-head dimension (None when it could not be derived, e.g. an export without KV-cache inputs).