aimet_onnx.experimental.spinquant

Top level APIs

aimet_onnx.experimental.spinquant.apply_spinquant(model, visual_model=None, embedding=None, *, enable_r1=True, enable_r2=False, enable_r3=False)[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. Analyzing the backbone (decoder block boundaries, role map, hidden size) and 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).

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).

Raises:

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

Return type:

None

Example (LLM):

apply_spinquant(model)
sim = QuantizationSimModel(model)
sim.compute_encodings(calibration_data)

Example (VLM):

embedding = torch.load("embedding.pth")   # torch.Tensor [vocab, hidden]
apply_spinquant(backbone_model, visual_model=visual_model, embedding=embedding)
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)