aimet_onnx.experimental.spinquant

Top level APIs

aimet_onnx.experimental.spinquant.apply_spinquant(backbone_sim, visual_sim=None, embedding=None, *, enable_r1=True, enable_r2=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 QuantizationSimModel(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).

  4. Rebuilding ORT sessions once per sim.

Must be called BEFORE sim.compute_encodings(). The rotation modifies float weight initializers; compute_encodings must run afterward to calibrate quantizer scales on the rotated weights.

Supported architectures:

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

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

Parameters:

  • backbone_sim (QuantizationSimModel) – QuantizationSimModel wrapping backbone.onnx.

  • visual_sim (Optional[QuantizationSimModel]) – Optional QuantizationSimModel wrapping visual.onnx (VLM only).

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

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

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

Example (VLM):

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