diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 85b5577494c..ea24e40f6ad 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -102,6 +102,7 @@ repos: # Instead, we should manually add the license header to those files *after* the original header. exclude: > (?x)^( + modelopt/torch/quantization/rotation/sgdg.py| modelopt/torch/quantization/utils/calib_utils.py| modelopt/onnx/quantization/operators.py| modelopt/onnx/quantization/ort_patching.py| diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 964fd8483fc..a85db0a4730 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,6 +8,7 @@ Changelog *Quantization* +- Add ``modelopt.torch.quantization.rotation`` — offline SpinQuant/QuaRot rotation folding (``fold_rotations``) and Cayley-SGD rotation learning (``learn_rotations``) of a global R1 plus per-layer R2 as a pre-quantization checkpoint transform for HF RMSNorm decoder LMs (Llama family, Qwen3). The rotated model is functionally identical (fp32 logits agree to ~3e-7) and remains a vanilla HF checkpoint, so every existing quant config, calibrator, exporter, and runtime works on it unchanged; online R3/R4 Hadamard transforms are out of scope. - Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Instead of setting the global scale from the largest per-block amax seen during calibration (plain ``max``, which leaves no room above it so any larger activation saturates), it anchors the scale to a low percentile of the per-block amax distribution, leaving the rest of the FP8 block-scale range as headroom: ``amax = max(rho * anchor, upper)``, where ``anchor`` and ``upper`` are the per-block amaxes at ``anchor_percentile`` (default 1) and ``upper_percentile`` (default 99.99; set to 100 to never clip calibration data), and ``rho`` (default 16384) is the headroom factor. Applies only to NVFP4 dynamic-block input quantizers; ``SequentialQuantizer`` activation quantizers raise. Weight scales are an orthogonal axis selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian``), so one recipe can combine a weight calibration with this activation policy in a single pass. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` with only the calibration algorithm swapped and exports a standard NVFP4 checkpoint. *Megatron Framework (M-LM / M-Bridge)* diff --git a/LICENSE b/LICENSE index a894d488493..1b809f6b2b0 100644 --- a/LICENSE +++ b/LICENSE @@ -244,6 +244,7 @@ the following copyright holders, licensed under the MIT License: Copyright (c) Andrei Panferov Copyright (c) Microsoft Corporation Copyright (c) 2020 EleutherAI + Copyright (c) 2020 Jun Li Copyright (c) 2020 Dan Hendrycks Copyright (c) 2023 Deep Cognition and Language Research (DeCLaRe) Lab Copyright (c) 2023 DeepSeek diff --git a/modelopt/torch/quantization/rotation/README.md b/modelopt/torch/quantization/rotation/README.md new file mode 100644 index 00000000000..ed080a8ef1c --- /dev/null +++ b/modelopt/torch/quantization/rotation/README.md @@ -0,0 +1,257 @@ +# Rotation Folding + Learning (SpinQuant/QuaRot R1 + R2) — Design + +Status: offline R1/R2 only (fold + Cayley-SGD learner); online R3/R4 are out of scope. +Module: `modelopt.torch.quantization.rotation` (`fold.py` + `learn.py` + `sgdg.py`). + +## Goal + +Apply SpinQuant/QuaRot-style rotations as a **pre-quantization checkpoint transform**: a +global orthogonal rotation R1 of the residual stream plus a per-layer head-space rotation R2 +on the v_proj → o_proj path are folded into the HF model weights in place. The rotated model +is functionally identical to the original (up to one float64 → original-dtype round-trip per +weight) but its activation/weight distributions are flatter, i.e. easier to quantize. The +transform is **orthogonal to qformat by construction**: it edits the checkpoint before +`mtq.quantize` runs, so every existing quant config, calibrator, exporter, and runtime works +unchanged on the rotated model. + +Mechanics are ported from a validated clean-room reference implementation +(fp32-equivalence validated on Qwen3-0.6B/1.7B — WikiText-2 PPL delta −0.0022% +(17.8337 → 17.8333), GSM8K checked; clean-room w.r.t. Meta's SpinQuant code). + +## Non-goals (explicit) + +- **No online transforms** — R3 (post-RoPE QK rotation) and R4 (down_proj activation + Hadamard) need runtime kernels; folding only their weight halves destroys the model. +- **No QuantAlgo / mode-registry / config-class integration** — one plain function. +- **No exporter or runtime changes** — the output is still a vanilla HF checkpoint. +- **No CLI** — callers script it. + +## API + +```python +from modelopt.torch.quantization.rotation import ( + fold_rotations, fold_seam_diags, learn_rotations, + QuantObjective, RotationSet, SGDG, + W4A4_G128_OBJECTIVE, INT8_DEFAULT_OBJECTIVE, W16A4_ASYM_R4G_OBJECTIVE, + SEAM_DIAG_LR, +) + +rotations = fold_rotations(model, mode="hadamard", seed=0, use_r2=True) +# model mutated in place (config.tie_word_embeddings forced False); +# rotations: {"R1": fp64 cpu [hidden, hidden], +# "model.layers.{i}.self_attn.R2": fp64 cpu [head_dim, head_dim]} + +rs = learn_rotations( + model, calib_loader, steps=150, lr=1.5, mode="hadamard", + objective_cfg=W4A4_G128_OBJECTIVE, # or any QuantObjective / None + seed=0, init_rotations=None, log_every=10, + teacher=None, kd_alpha=0.5, kd_temp=2.0, # optional KD objective +) # -> RotationSet: .rotations (R.bin convention), .history, .meta, + # .seam_diags (transform-QAT only, else None) + +fold_rotations(model, R1=rs.R1, R2=rs.R2) # bake a checkpoint +fold_seam_diags(model, rs.seam_diags, smax=256) # transform-QAT: bake seam scales +``` + +Pipeline order (load-bearing): untie tied embeddings with a real clone → seed the global +torch RNG → fuse RMSNorm gains into downstream linears (fused norms become exactly ones; +pure RMSNorm commutes with an orthogonal rotation of its input) → apply R1 → apply R2. +All math in float64, cast back to the original dtype. + +## Reader/writer orientation + +| Weight | Role vs residual stream | Transform | +|---|---|---| +| `embed_tokens` | writer (rows are stream vectors) | `E ← E @ R1` | +| `q/k/v_proj`, `gate/up_proj` | reader (input side) | `W ← W @ R1` | +| `o_proj`, `down_proj` | writer (output side) | `W ← R1ᵀ @ W`, `b ← R1ᵀ b` | +| `lm_head` (after final-norm fusion) | reader | `W ← W @ R1` | +| `v_proj` (R2) | output side, per **KV** head | `W_h ← R2ᵀ @ W_h` (row blocks) | +| `o_proj` (R2) | input side, per **Q** head | `W[:, h·d:(h+1)·d] ← … @ R2` (col blocks) | +| `q_norm`/`k_norm` (Qwen3) | per-head, post-q/k_proj head space | **never fused, never rotated** | + +One R2 is shared by all heads of a layer, which is what keeps GQA/`repeat_kv` exact. + +## Arch-mapping registry + +`_ARCH_REGISTRY` is one small dict keyed on the model **class name** (`LlamaForCausalLM`, +`Qwen3ForCausalLM`), with fields: `has_qk_norm`, `head_dim` (callable on the config — must +prefer `config.head_dim`; `hidden_size // num_attention_heads` is wrong for Qwen3-0.6B: 64 +vs the true 128), and `norm_edges` (RMSNorm → downstream-linear fusion edges). Any other +class raises `NotImplementedError`. Adding a standard-layout HF decoder +(`model.model.{embed_tokens,layers,norm}` + `model.lm_head`) is one dict entry. + +## Equivalence gates + +1. **Generation gate** — every R asserts `max |R Rᵀ − I| < 1e-10` in fp64 at build time. +2. **In-function post-conditions** — fused norms exactly ones; Qwen3 q/k_norm bitwise + untouched; every parameter shape unchanged. +3. **Unit gate** — fp32 logits before vs after fold agree to `atol=1e-4` (measured ~3e-7); + `tests/unit/torch/quantization/test_rotation_fold.py`. +4. **Model gate (external, before any PTQ)** — WikiText-2 PPL of the rotated fp checkpoint + must match the original to noise. + +## External-matrix fold path (learned rotations enter here) + +`fold_rotations(model, R1=..., R2=...)` folds externally supplied matrices through the +exact same pipeline (untie → fuse → rotate). The seed path is unchanged and bitwise +reproducible; the external path skips the RNG entirely (global RNG state untouched) and +gates each matrix at `max |R Rᵀ − I| < 1e-4` — the *trained-rotation* deployability +tolerance (fp32 Cayley iterates drift ~1e-7/step; 150-step reference runs measure ~5e-5) +instead of the 1e-10 fresh-draw gate. `R2` accepts a layer-ordered sequence, an +int-keyed dict, or the returned-dict key convention (`model.layers.{i}.self_attn.R2`). +Unit gate: matrices returned by the seed path, fed back through `R1=`/`R2=` on an +identical model, reproduce every parameter bitwise. + +## Learned rotations — `learn.py` + +`learn_rotations(model, calib_loader, steps=150, lr=1.5, mode="hadamard", +objective_cfg=..., teacher=None, ...) -> RotationSet` learns the same offline pair the +fold applies (SpinQuant, arXiv:2405.16406): + +- **What is learned**: the global residual-stream rotation **R1 `[hidden, hidden]`** plus + one per-layer head-space rotation **R2 `[head_dim, head_dim]`** on the v_proj → o_proj + path, as fp32 parameters on the **Stiefel manifold**, updated by **Cayley SGD** (the + `SGDG` optimizer of Li et al., MIT-licensed, ported self-contained into `sgdg.py` — no new + deps; the original momentum dead-store quirk is reproduced and documented, so momentum is inert + exactly as in the official trainer). +- **Objective**: next-token **CE of the fake-quantized rotated model** on calibration + text. Each step assembles every rotated effective weight out-of-place per fold.py's + orientation table (readers `W @ R1`, writers `R1ᵀ @ W`, embed/lm_head included, v/o R2 + block mechanics), applies the objective's weight fake-quant with a straight-through + estimator, reparametrizes the *frozen* model with those tensors + (`torch.nn.utils.stateless._reparametrize_module`, spanning forward and backward), and + fake-quantizes activations via pre-hooks on the 7×n_layers target linears. Gradients + reach only R1/R2 (asserted at step 0; plus the seam-diag leaves under transform-QAT). + Cosine lr decay over `steps`. +- **KD objective (optional)**: `teacher=` a frozen reference model (typically a + bf16 copy of the SAME checkpoint) switches the loss to + `(1−kd_alpha)·CE + kd_alpha·kd_temp²·KL(student ‖ teacher)`, teacher logits under + `no_grad` on the same batch. The teacher is never reparametrized/fused/modified + (unit-asserted). `teacher=None` (default) is bitwise the plain-CE trainer. + `meta["kd"]` records `{alpha, temp}` when active. +- **Transform-QAT (optional, `QuantObjective.learn_seam_diag=True`)**: jointly learns + OSTQuant-style per-input-channel diagonal scales at the two rotation-SURVIVING seams + (down_proj input `[intermediate]`, o_proj input `[n_kv·head_dim]`, GQA-exact) as + `log s` leaves (init 0 = identity; no extra RNG consumed — the R trajectory stream is + unchanged). They live in a separate plain-Adam group at `SEAM_DIAG_LR` (1e-2, same + cosine schedule — never the SGDG stiefel group), are applied in the effective-weight + assembly with the SmoothQuant-style prefold structure (prefold-inside / rotation-outside), and + export as `RotationSet.seam_diags` (fp64, positive by construction). Bake with + `fold_seam_diags(model, seam_diags, smax=256)` (fp64 exact identities; the smax + ceiling is the fp16-subnormal guard). `save()`/`load()` round-trips them under a + reserved key; legacy flat R.bins load with `seam_diags=None`. +- **Init**: random(-sign) **Hadamard** (or Haar `"random"`) drawn with the SAME seeded + global-RNG draw order as `fold_rotations` — `steps=0` returns bitwise the fold's seed + draws (unit-gated), so trained and random rotations share one provenance contract. + Warm starts enter via `init_rotations` (R.bin key convention, ortho-gated). +- **Model side effects** (identical to fold's pre-rotation steps, so learn-then-fold on + one object or fold-on-a-fresh-copy are both valid): untie with a real clone, RMSNorm + gains fused (idempotent), params frozen. Weights themselves are never rewritten — + rotated weights exist only inside the per-step reparametrization. Qwen3 q/k_norm + bitwise untouched (asserted; head-space, post-projection — reused arch spec). +- **Final retraction (load-bearing)**: raw fp32 Cayley iterates drift off the manifold, + and the max-entry residual is **basis-dependent**: field measurements on 150-step R1s + give `max|RᵀR−I| ≈ 4.6e-5 / 8.0e-5` (the form the step audit reports) but + `max|RRᵀ−I| ≈ 7.3e-4 / 1.6e-3` — 10–20× larger — and the fold orientation consumes + the `RRᵀ` form (reader/writer seams compose to `x R1 R1ᵀ Wᵀ`). Before returning, + every trained matrix is polar-projected to the nearest orthogonal matrix + (`R = UΣVᵀ → UVᵀ`, float64; per-entry move ≈ drift/2, far below one bf16 ulp; the + same retraction semantics SGDG applies stochastically during training, applied once + deterministically at the end). Post-retraction residual ~1e-14 both forms; raw + drift + projection distance recorded in `meta["final_retraction"]`. The predecessor + consumer of R.bin files handled the same asymmetry by widening + its gate to 1e-3; the module closes the drift instead of widening the gate. +- **Output**: `RotationSet` — float64 CPU dict in the fold/R.bin key convention plus + per-step `history` and `meta` (final static-A8 amax and the retraction log included + when applicable). `save()`/`load()` round-trips the flat dict bitwise; `load()` + refuses off-manifold matrices (`orthogonalize=True` retracts raw legacy R.bins). + Feed `fold_rotations(model, R1=rs.R1, R2=rs.R2)` to bake a checkpoint. + +Architecture knowledge is REUSED from `fold.py` (`_ARCH_REGISTRY`, norm edges, head_dim +resolution, q/k_norm exclusion, untie handling) — defined once, imported by the learner. + +### Pluggable fake-quant objectives (`QuantObjective`) + +Default numerics: symmetric integer QDQ, ModelOpt max-calibration semantics — +`s = amax/(2^{b-1}−1)` (clamped 1e-12), round-half-even, clamp `[−2^{b-1}, 2^{b-1}−1]`; +STE backward. Two paper-protocol extensions: + +- **`a_asym=True`** — per-token dynamic ASYMMETRIC min-max affine activation QDQ, + matching the official SpinQuant `ActQuantizer` (`sym=False`, `clip_ratio=1`) + bit-for-bit: zero-inclusive token range, all-zero-token fallback to `[−1, 1]`, + `scale=(max−min)/(2^b−1)`, `zp=round(−min/scale)` (unit-tested against a line-for-line + reference transcription). Per-token-dynamic mode only. +- **`r4_in_graph=True`** — the online R4 down_proj Hadamard placed in the TRAINING + graph only (input hook `x @ H` before act-QDQ + effective down_proj columns `@ H`; a + functional-identity pair that only the quantizers see). The deployed fold never sees + H (unit-asserted). Power-of-2 seam dims only (`_walsh_hadamard`; Llama-3.2 8192 ✓, + Qwen3's 6144 needs the unimplemented had-K composition). + +| preset | weights | activations | role | +|---|---|---|---| +| `W4A4_G128_OBJECTIVE` | int4 sym per-group g128 (in-dim) | int4 sym **per-token dynamic** | SpinQuant-paper-style W4A4; reference-trainer comparison axis | +| `INT8_DEFAULT_OBJECTIVE` | int8 sym **per-out-channel** | int8 sym **per-tensor STATIC** | the axes of ModelOpt `INT8_DEFAULT_CFG` — the deployment cell where random rotations barely help | +| `W16A4_ASYM_R4G_OBJECTIVE` | none (W16 in the loss) | int4 **asym per-token dynamic** + R4-in-graph | the official GPTQ-deploy training objective ("Cayley on 16-4-KV", paper Table 3) — kept for ablation; measured HARMFUL for R1R2-only deployment (objective-lever anchor below) | +| custom `QuantObjective(...)` | any bits, per-group or per-channel | per-token dynamic (sym/asym) / per-tensor static / off; `learn_seam_diag`, `r4_in_graph` | e.g. exact reference-trainer replica = `w_bits=4, w_group=None, a_bits=4, per_token_dynamic`; recommended W4A4 objective = `w4a4_g128 + a_asym=True`, built via `QuantObjective(name="w4a4_g128_asym", w_bits=4, w_group=128, a_bits=4, a_asym=True)` — NOT the shipped default: `W4A4_G128_OBJECTIVE` stays symmetric for paper-protocol comparability | + +Coverage: the 7 per-layer projections only; embeddings/lm_head never quantized (matches +`INT8_DEFAULT_CFG`'s `*lm_head*` exclusion and the deployed ModelOpt cells). Static-A8 +scale scope is `a_static_scope="batch"` by DEFAULT — a fresh amax per calib batch, the +stationary surrogate that tracks the moving rotation (post-hoc max calibration of the +folded ckpt is what deployment sees anyway). The literal running-max semantics +(`"run"`) is kept for ablation only: measured NON-stationary under a moving rotation +(loss drifts up as stale outliers pin the scale). + +### Lineage diff — official SpinQuant vs. our internal reference trainer vs. this module + +The middle column is an internal reference reimplementation (not shipped); it is kept in +the table because it is where the numerics of this module were first validated. + +| axis | official `optimize_rotation.py` | internal reference trainer (not shipped) | `learn.py` (this module) | +|---|---|---|---| +| trained params | R1 + per-layer R2, fp32 `RotateModule`s | R1 + per-layer R2, fp32 `nn.Parameter`s | same as reference | +| R2 size derivation | `hidden_size // num_attention_heads` (**wrong for Qwen3-0.6B**: 64 vs true 128) | `config.head_dim` (asserted) | fold.py registry (`config.head_dim` preferred/asserted) | +| model plumbing | forked `modeling_llama_quant.py` with online rotation modules + QuantizeLinear wrappers (Llama only) | stock HF model, frozen; per-step effective-weight assembly + reparametrize | same as reference, arch registry Llama+Qwen3 | +| quantizer in the loss | their QuantLinear: asym per-token A4 (`--a_asym`), clip-searched weights (`--w_clip`); trained vs **activation-quant-only** network when pairing with GPTQ | v1: per-out-channel sym W4 (max scale, ±7 clamp) + per-token dynamic sym A4 STE hooks; v2: clip-searched W + asym A4 | pluggable `QuantObjective` (table above); sym, ModelOpt max-scale numerics, `[−2^{b−1}, 2^{b−1}−1]` | +| calib data | WikiText-2 train, seq 2048, 800 samples | GSM8K-train blocks, 512 tokens (an earlier variant used WikiText-2) | caller-supplied loader (any re-iterable of input_ids) | +| budget | 100 steps, effective bs 8 (1×8 GPUs), lr 1.5 | 150 steps, bs 4, seq 512, lr 1.5 | defaults steps=150, lr=1.5; caller-set | +| lr schedule | cosine (HF Trainer `--lr_scheduler_type cosine`) | cosine (explicit) | cosine (explicit, same formula) | +| optimizer | SGDG stiefel=True | SGDG port (stiefel branch verified bitwise vs official; momentum dead-store documented) | same port, self-contained | +| norm fusion / untie | `prepare_model` fuses norms; lm_head cloned for 3.2 | norm-fusion helper shared with the reference fold path; untie clone | reused `fold.py` helpers; untie clone | +| dtype discipline | bf16 model, fp32 R | bf16 model, fp32 R, TF32 force-off, bf16-grid store points mirroring the offline fold (bitwise store-point anchor) | model dtype preserved (fp32 CPU tests / bf16 GPU), fp32 R; no intermediate-grid mirroring (functional, not bitwise, contract) | +| output | `R.bin` (fp32 state-dict values) | `R.bin` fp64 + loss.jsonl + config.json | `RotationSet` (fp64, R.bin-compatible `save()`, history+meta in-object) | +| paper/no-online-Hadamard anchor | ~~v1 Table 3: R1-only 9.6~~ **superseded by v4 (ICLR'25) Table 8**: Llama-3.2-1B W4A4KV16 "SpinQuant no had" (= learned R1R2-only, GPTQ+w_clip+a_asym) wiki **48.4** vs fp **13.4** = **3.61×**; "SpinQuant had" (+online R3/R4) 15.3 = 1.14× | reference-trainer Llama-3.2-1B W4A4: 23.32 (windowed protocol, bf16 8.59) | measured (module training runs + official-harness alignment matrix) | + +### Anchor correction — paper v4 (ICLR'25) no-had W4A4 band + +The v1-based anchor above was misleading: arXiv 2405.16406 **v4** shows that a learned +R1R2-only ("no had") rotation at **W4A4 is NOT close to the fp baseline on any model** — +Llama-3.2-1B: 48.4 vs fp 13.4 (3.61×); LLaMA-2-7B: 9.2 vs 5.5 (1.67×); L3-8B: 18.6 vs 6.1 +(3.05×). "No-had ≈ fp" is the paper's **W4A8** claim (1B: 15.3 vs 13.4). Closing the W4A4 gap +requires the online R3/R4 Hadamard ("had" scheme). Cross-checked on the official harness +(external alignment study, Llama-3.2-1B): +this module's learned R.bin dropped into official `ptq.py` gives no-had 61.2 (RTN+clip+asym) +vs random-no-had 109.2 vs none 256.1, and 18.8 with online R4 — same structure as the paper. +Budget note: paper Table 11 shows rotation quality saturates at 100 iters / 128 samples, so +the module's default 150 steps is not the binding factor at W4A4; the eval recipe +(GPTQ/clip/asym + online R4) is. + +### Objective-lever anchor — Llama-3.2-1B W4A4 no-had + +Measured on Llama-3.2-1B W4A4KV16, official released harness, no-had (R1R2-only) +deployment, GPTQ+clip+asym eval (external measurement campaign): + +| training objective | no-had GPTQ PPL | verdict | +|---|---|---| +| `w4a4_g128` (sym, the shipped default — kept for paper-protocol comparability) | 57.27 | regression band | +| **`w4a4_g128 + a_asym=True`** | **48.89 — within 1% of the paper's 48.4** | **recommended W4A4 objective** (not the shipped default; see the objective table) | +| `+ r4_in_graph` (W4 in loss) | 59.55 | HARMFUL — trains a grid deployment never uses | +| `W16A4_ASYM_R4G` (Table-3 objective) | 93.72 | HARMFUL — replicates the official-trainer arm4 negative (cross-trainer consistency) | + +The residual vs the paper's normalized multiple (5.01× vs 3.61× own-fp) equals the +internal-vs-released fp-anchor discrepancy (13.4/9.7611 = 1.37×) documented in the anchor +correction above — on the released harness, parity within ~1% is the achievable ceiling. + +Online transforms require exporter + runtime kernel support (TRT-LLM) and remain a +separate track; `learn.py` trains only the foldable pair R1/R2. diff --git a/modelopt/torch/quantization/rotation/__init__.py b/modelopt/torch/quantization/rotation/__init__.py new file mode 100644 index 00000000000..ec4e1c4c4b9 --- /dev/null +++ b/modelopt/torch/quantization/rotation/__init__.py @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Rotation folding + learning (SpinQuant/QuaRot R1 + R2) as pre-quantization transforms.""" + +from .fold import * +from .learn import * diff --git a/modelopt/torch/quantization/rotation/fold.py b/modelopt/torch/quantization/rotation/fold.py new file mode 100644 index 00000000000..89ff032ffc0 --- /dev/null +++ b/modelopt/torch/quantization/rotation/fold.py @@ -0,0 +1,615 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Offline SpinQuant/QuaRot rotation folding (R1 + R2) for HF RMSNorm decoder LMs. + +:meth:`fold_rotations` rewrites the weights of a supported HuggingFace causal LM in place so +that a global orthogonal rotation R1 of the residual stream (and optionally a per-layer +head-space rotation R2 on the v_proj -> o_proj path) is folded into the checkpoint. The +transform is a functional identity up to one float64 -> original-dtype round-trip per weight, +and is applied *before* quantization: rotated activation/weight distributions are flatter +(fewer outliers) and therefore easier to quantize (SpinQuant, QuaRot). + +Only the offline rotations are applied here. SpinQuant's online transforms (R3 post-RoPE QK +rotation, R4 down_proj activation Hadamard) require runtime kernel support and are out of +scope — see README.md in this directory. +""" + +import math +from typing import Any + +import torch +import torch.nn as nn + +__all__ = ["fold_rotations", "fold_seam_diags"] + +# -------------------------------------------------------------------------------------- +# Random orthogonal matrix generation (random-sign Hadamard D @ H / sqrt(n), or Haar QR) +# -------------------------------------------------------------------------------------- + +# SpinQuant's priority order of hard-coded had-K block sizes; the subset with q = K - 1 +# prime and q ≡ 3 (mod 4) is regenerated via Paley construction I. Any valid Hadamard +# decomposition yields an orthonormal matrix, which is all that matters for fresh random +# rotations (bitwise parity with SpinQuant's hard-coded had-K constants is irrelevant). +_HADK_PRIORITY = [172, 156, 140, 108, 60, 52, 36, 28, 44, 40, 20, 12] +_HADK_SUPPORTED = {12, 20, 44, 60, 108, 140} + + +def _is_pow2(n: int) -> bool: + return n > 0 and (n & (n - 1)) == 0 + + +def _paley_hadamard(K: int) -> torch.Tensor: + """Hadamard matrix of order K = q + 1 via Paley construction I (q prime, q ≡ 3 mod 4).""" + q = K - 1 + residues = {(i * i) % q for i in range(1, q)} + + def chi(a: int) -> float: + a %= q + return 0.0 if a == 0 else (1.0 if a in residues else -1.0) + + S = torch.zeros(K, K, dtype=torch.float64) + S[0, 1:] = 1.0 + S[1:, 0] = -1.0 + for i in range(q): + for j in range(q): + if i != j: + S[1 + i, 1 + j] = chi(i - j) + H = S + torch.eye(K, dtype=torch.float64) + assert torch.allclose(H @ H.T, K * torch.eye(K, dtype=torch.float64)), ( + f"Paley construction failed for K={K}" + ) + return H + + +def _get_hadK(n: int) -> tuple[torch.Tensor | None, int]: + """Return ``(hadK matrix or None, K)`` such that ``n = 2^k * K``.""" + if _is_pow2(n): + return None, 1 + for K in _HADK_PRIORITY: + if n % K == 0 and _is_pow2(n // K) and K in _HADK_SUPPORTED: + return _paley_hadamard(K), K + raise ValueError( + f"size {n} is not 2^k or 2^k*K for a regenerable K {sorted(_HADK_SUPPORTED)}; " + 'use mode="random" instead' + ) + + +def _matmul_hadU(X: torch.Tensor) -> torch.Tensor: + """Fast Walsh-Hadamard transform over the last dim of X. + + Recurses down to block size K, then one ``hadK @ blocks`` matmul, then ``/ sqrt(n)`` + (the normalization is applied exactly once). + """ + n = X.shape[-1] + hadK, K = _get_hadK(n) + inp = X.clone().view(-1, n, 1) + out = inp.clone() + while inp.shape[1] > K: + inp = inp.view(inp.shape[0], inp.shape[1] // 2, 2, inp.shape[2]) + out = out.view(inp.shape) + out[:, :, 0, :] = inp[:, :, 0, :] + inp[:, :, 1, :] + out[:, :, 1, :] = inp[:, :, 0, :] - inp[:, :, 1, :] + out = out.view(inp.shape[0], inp.shape[1], -1) + inp, out = out, inp + del out + if K > 1: + assert hadK is not None + inp = hadK.view(1, K, K).to(inp) @ inp + return inp.view(X.shape) / math.sqrt(n) + + +def _random_hadamard_matrix(size: int) -> torch.Tensor: + """``D @ H / sqrt(n)``: random-sign Hadamard. + + Signs come from the global CPU torch RNG (seeded once in :meth:`fold_rotations`; the draw + order is load-bearing for seed reproducibility). + """ + signs = torch.randint(0, 2, (size,)) + Q = torch.diag(signs.to(torch.float64) * 2 - 1) + return _matmul_hadU(Q) + + +def _random_orthogonal_matrix(size: int) -> torch.Tensor: + """QR of a float64 randn with sign fix -> Haar-uniform orthogonal Q.""" + m = torch.randn(size, size, dtype=torch.float64) + q, r = torch.linalg.qr(m) + q *= torch.sign(torch.diag(r)).unsqueeze(0) + return q + + +def _get_orthogonal_matrix(size: int, mode: str) -> torch.Tensor: + if mode == "hadamard": + R = _random_hadamard_matrix(size) + elif mode == "random": + R = _random_orthogonal_matrix(size) + else: + raise ValueError(f"unknown rotation mode: {mode!r} (expected 'hadamard' or 'random')") + err = (R @ R.T - torch.eye(size, dtype=torch.float64)).abs().max().item() + assert err < 1e-10, f"generated matrix not orthogonal (max |R R^T - I| = {err:.3e})" + return R + + +# Externally supplied matrices are gated at 1e-4 (vs 1e-10 for fresh draws). The gate +# checks the ``R R^T`` form because that is what the fold orientation consumes (the +# reader/writer seams compose to ``x R1 R1^T W^T``) — and for near-orthogonal trained +# matrices the max-entry residual of ``R R^T - I`` is 10-20x larger than the ``R^T R`` +# form the trainer's step audit reports (basis-dependent; measured raw 150-step R1: +# 5e-5 vs 1e-3). learn_rotations closes this with a final polar retraction (~1e-14 both +# forms), so learned RotationSets pass with huge headroom; raw legacy R.bins must be +# loaded with RotationSet.load(..., orthogonalize=True) first. +_EXTERNAL_ORTHO_TOL = 1e-4 + + +def _as_external_rotation(mat, size: int, name: str) -> torch.Tensor: + """Validate one externally supplied rotation: shape [size, size], orthonormal. + + Orthonormality is gated at :data:`_EXTERNAL_ORTHO_TOL`. Returns a float64 CPU copy (all + fold math is float64). The copy is unconditional: ``as_tensor``/``to``/``cpu`` are + no-ops for an already-float64 CPU input, and without it the audited matrix would alias + the caller's buffer and could change after the gate passed. + """ + R = torch.as_tensor(mat).detach().to(torch.float64).cpu().clone() + if R.shape != (size, size): + raise ValueError(f"{name}: expected shape {(size, size)}, got {tuple(R.shape)}") + err = (R @ R.T - torch.eye(size, dtype=torch.float64)).abs().max().item() + if err >= _EXTERNAL_ORTHO_TOL: + raise ValueError( + f"{name}: not orthogonal (max |R R^T - I| = {err:.3e} >= " + f"{_EXTERNAL_ORTHO_TOL}) — refusing to fold a non-orthogonal rotation" + ) + return R + + +def _normalize_external_r2(R2, n_layers: int, head_dim: int) -> list[torch.Tensor]: + """Normalize the accepted per-layer R2 formats into a validated list ordered by layer. + + Accepts a sequence ordered by layer, a dict keyed by layer index, or a dict in the R.bin + key convention (``model.layers.{i}.self_attn.R2`` — the format fold_rotations returns and + RotationSet.R2 provides). + """ + if isinstance(R2, dict): + by_idx = {} + for k, v in R2.items(): + if isinstance(k, int): + idx = k + else: + try: + idx = int(str(k).split("model.layers.")[1].split(".")[0]) + except (IndexError, ValueError): + raise ValueError( + f"R2 dict key {k!r} not understood (want an int layer index or " + "'model.layers.{i}.self_attn.R2')" + ) from None + if idx in by_idx: + raise ValueError( + f"R2 dict names layer {idx} more than once (e.g. the int key {idx} and " + f"'model.layers.{idx}.self_attn.R2'); refusing to guess which matrix wins" + ) + by_idx[idx] = v + if sorted(by_idx) != list(range(n_layers)): + raise ValueError( + f"R2 must cover every layer 0..{n_layers - 1}; got indices {sorted(by_idx)}" + ) + mats = [by_idx[i] for i in range(n_layers)] + else: + mats = list(R2) + if len(mats) != n_layers: + raise ValueError(f"R2 has {len(mats)} matrices for {n_layers} layers") + return [_as_external_rotation(m, head_dim, f"R2[layer {i}]") for i, m in enumerate(mats)] + + +# -------------------------------------------------------------------------------------- +# Architecture mapping registry +# -------------------------------------------------------------------------------------- + + +def _llama_head_dim(cfg) -> int: + return getattr(cfg, "head_dim", None) or cfg.hidden_size // cfg.num_attention_heads + + +def _qwen3_head_dim(cfg) -> int: + head_dim = getattr(cfg, "head_dim", None) + assert head_dim is not None, ( + "config.head_dim missing — refusing to fall back to hidden_size // num_attention_heads " + "(wrong for Qwen3-0.6B: that gives 64 but the true head_dim is 128)" + ) + return head_dim + + +# RMSNorm -> downstream-linear fusion edges, relative to one decoder layer. The final-norm -> +# lm_head edge is handled explicitly in fold_rotations. Per-head q_norm/k_norm (Qwen3) act in +# post-q/k_proj head space, not on the residual stream: they are NEVER fused, NEVER rotated. +_NORM_EDGES = ( + ("input_layernorm", ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj")), + ("post_attention_layernorm", ("mlp.gate_proj", "mlp.up_proj")), +) + +# Keyed on the model class name; every entry must follow the standard HF decoder layout +# (model.model.{embed_tokens,layers,norm} + model.lm_head). +_ARCH_REGISTRY: dict[str, dict[str, Any]] = { + "LlamaForCausalLM": { + "has_qk_norm": False, + "head_dim": _llama_head_dim, + "norm_edges": _NORM_EDGES, + }, + "Qwen3ForCausalLM": { + "has_qk_norm": True, + "head_dim": _qwen3_head_dim, + "norm_edges": _NORM_EDGES, + }, +} + + +# -------------------------------------------------------------------------------------- +# Norm fusion and rotation application (all math in float64, cast back to original dtype) +# -------------------------------------------------------------------------------------- + + +def _fuse_norm_into_linears(norm: nn.Module, linears: list[nn.Module]) -> None: + """Fold the RMSNorm gain into the input columns of each downstream linear. + + ``W <- (W.double() * gamma.double()).to(orig_dtype)``; the norm weight becomes ones. + RMSNorm only — a norm bias (LayerNorm) is not supported. + """ + assert getattr(norm, "bias", None) is None, ( + "unexpected bias on norm — LayerNorm fusion is not ported (RMSNorm only)" + ) + gamma = norm.weight.data.to(torch.float64) + for lin in linears: + w = lin.weight + assert w.shape[1] == gamma.numel(), ( + f"fuse mismatch: W in_dim {w.shape[1]} vs gamma {gamma.numel()}" + ) + w.data = (w.data.to(torch.float64) * gamma).to(w.dtype) + norm.weight.data = torch.ones_like(norm.weight.data) + + +def _rotate_input_cols(module: nn.Module, R1: torch.Tensor, row_chunk: int = 32768) -> None: + """Reader/embedding rotation: ``W <- W @ R1``. + + Rows are chunked to bound the float64 temporary for vocab-sized matrices. Bias is never + touched (input-side rotation). + """ + w = module.weight + R = R1.to(w.device) + for s in range(0, w.shape[0], row_chunk): + w.data[s : s + row_chunk] = (w.data[s : s + row_chunk].to(torch.float64) @ R).to(w.dtype) + + +def _rotate_output_rows(linear: nn.Module, R1: torch.Tensor) -> None: + """Writer rotation (o_proj / down_proj): ``W <- R1^T @ W``; bias ``b <- R1^T b``.""" + w = linear.weight + R = R1.to(w.device) + w.data = (R.T @ w.data.to(torch.float64)).to(w.dtype) + if linear.bias is not None: + b = linear.bias + b.data = (R.T @ b.data.to(torch.float64)).to(b.dtype) + + +def _rotate_v_proj_r2(linear: nn.Module, R2: torch.Tensor, head_dim: int) -> None: + """Per-KV-head block of rows: ``W_h <- R2^T @ W_h``. + + Implemented as transpose, reshape to ``[in, out//d, d]``, right-multiply by R2, then + reshape/transpose back. + """ + assert linear.bias is None, ( + "v_proj has a bias — R2 would silently break equivalence (bias must be rotated per head)" + ) + w = linear.weight + out_f, in_f = w.shape + assert out_f % head_dim == 0, f"v_proj out {out_f} not divisible by head_dim {head_dim}" + R = R2.to(w.device) + Wt = w.data.to(torch.float64).t() + Wt = (Wt.reshape(in_f, out_f // head_dim, head_dim) @ R).reshape(in_f, out_f) + w.data = Wt.t().contiguous().to(w.dtype) + + +def _rotate_o_proj_r2(linear: nn.Module, R2: torch.Tensor, head_dim: int) -> None: + """Per-Q-head block of columns: ``W[:, h*d:(h+1)*d] <- W[:, h*d:(h+1)*d] @ R2``. + + The block count derives from the tensor shape (in_features = num_q_heads * head_dim), so + GQA is handled automatically. Input-side w.r.t. v: bias untouched. + """ + w = linear.weight + out_f, in_f = w.shape + assert in_f % head_dim == 0, f"o_proj in {in_f} not divisible by head_dim {head_dim}" + R = R2.to(w.device) + W = w.data.to(torch.float64).reshape(out_f, in_f // head_dim, head_dim) @ R + w.data = W.reshape(out_f, in_f).to(w.dtype) + + +# -------------------------------------------------------------------------------------- +# Public API +# -------------------------------------------------------------------------------------- + + +def fold_rotations( + model: nn.Module, + mode: str = "hadamard", + seed: int = 0, + use_r2: bool = True, + R1: torch.Tensor | None = None, + R2=None, +) -> dict[str, torch.Tensor]: + """Fold offline SpinQuant/QuaRot rotations (R1 + per-layer R2) into ``model`` in place. + + Pipeline (order is load-bearing): untie tied embeddings with a real clone -> seed the + global torch RNG -> fuse RMSNorm gains into downstream linears -> apply R1 (readers + ``W @ R1``, writers ``R1^T @ W``; embed_tokens and lm_head included) -> apply one shared + per-layer R2 on the v_proj -> o_proj head space. All rotation math runs in float64 and + is cast back to the original weight dtype, so the model output is unchanged up to that + round-trip. + + Rotation source — two mutually exclusive paths: + + * **Seed path (default)**: ``R1``/``R2`` omitted; matrices are drawn from the seeded + global RNG per ``mode`` (byte-identical for equal seeds; unchanged legacy behavior). + * **External path**: pass ``R1`` (and ``R2`` when ``use_r2=True``) explicitly — e.g. + matrices learned by :meth:`~modelopt.torch.quantization.rotation.learn_rotations` — + and they are folded through this same validated path. External matrices are gated at + orthonormality :data:`_EXTERNAL_ORTHO_TOL` (1e-4, the trained-rotation deployability + tolerance) instead of the 1e-10 fresh-draw gate; ``mode``/``seed`` are ignored. + + Args: + model: HuggingFace causal LM whose class is registered (currently + ``LlamaForCausalLM`` and ``Qwen3ForCausalLM``). Modified in place; + ``config.tie_word_embeddings`` is set to False. + mode: ``"hadamard"`` (random-sign Hadamard, ``D @ H / sqrt(n)``) or ``"random"`` + (Haar-uniform orthogonal via QR). Seed path only. + seed: Seed for the global CPU torch RNG; all rotation matrices draw from it in a + fixed order, so equal seeds give byte-identical rotations. Seed path only. + use_r2: Also apply the per-layer head-space rotation R2 (one shared matrix per + layer across all heads — required for GQA / repeat_kv correctness). + R1: Optional external global rotation ``[hidden, hidden]``. + R2: Optional external per-layer head-space rotations: a sequence ordered by layer, + a dict keyed by layer index, or a dict in the returned key convention + (``model.layers.{i}.self_attn.R2``). Requires ``R1`` and ``use_r2=True``. + + Returns: + The applied rotations as float64 CPU tensors, keyed ``"R1"`` plus + ``"model.layers.{i}.self_attn.R2"`` (SpinQuant optimized-checkpoint convention). + + Raises: + NotImplementedError: If the model architecture is not in the registry. + ValueError: If external matrices are inconsistent (shape/count/orthogonality, R2 + without R1, or a missing R2 with ``use_r2=True``). + """ + arch = type(model).__name__ + if arch not in _ARCH_REGISTRY: + raise NotImplementedError( + f"fold_rotations: unsupported architecture {arch!r}; " + f"supported: {sorted(_ARCH_REGISTRY)}" + ) + spec = _ARCH_REGISTRY[arch] + + decoder = model.model + layers = decoder.layers + embed = decoder.embed_tokens + head_dim = spec["head_dim"](model.config) + + # External-matrix path validation (learned rotations enter here). + external = R1 is not None or R2 is not None + r1_ext = r2_ext = None + if external: + if R1 is None: + raise ValueError("external rotations: R1 is required when R2 is given") + if use_r2 and R2 is None: + raise ValueError( + "external rotations: use_r2=True needs R2 matrices (or pass use_r2=False)" + ) + if not use_r2 and R2 is not None: + raise ValueError("external rotations: R2 given but use_r2=False") + r1_ext = _as_external_rotation(R1, model.config.hidden_size, "R1") + if use_r2: + r2_ext = _normalize_external_r2(R2, len(layers), head_dim) + + # 1. Untie embeddings with a real clone: lm_head (reader) and embed_tokens (writer) + # diverge below because only lm_head absorbs the final-norm gain. + if model.lm_head.weight.data_ptr() == embed.weight.data_ptr(): + model.lm_head.weight = nn.Parameter( + embed.weight.data.clone(), requires_grad=embed.weight.requires_grad + ) + model.config.tie_word_embeddings = False + assert model.lm_head.weight.data_ptr() != embed.weight.data_ptr(), "untie failed" + + # Snapshots for post-condition checks (after untie: named_parameters() deduplicates a + # tied lm_head.weight, so the snapshot would otherwise miss it). + shapes_before = {n: tuple(p.shape) for n, p in model.named_parameters()} + qk_norm_before = {} + if spec["has_qk_norm"]: + for idx, layer in enumerate(layers): + qk_norm_before[f"{idx}.q_norm"] = layer.self_attn.q_norm.weight.data.clone() + qk_norm_before[f"{idx}.k_norm"] = layer.self_attn.k_norm.weight.data.clone() + + # 2. Seed the global CPU RNG (every rotation matrix draws from it, in a fixed order). + # External path: nothing is drawn, so the global RNG state is left untouched. + if not external: + torch.manual_seed(seed) + + # 3. Fuse RMSNorm gains into downstream linears (fused norms become exactly ones). + for layer in layers: + for norm_name, linear_names in spec["norm_edges"]: + _fuse_norm_into_linears( + layer.get_submodule(norm_name), [layer.get_submodule(n) for n in linear_names] + ) + _fuse_norm_into_linears(decoder.norm, [model.lm_head]) + + # 4. Rotate: R1 over the residual stream, then per-layer R2 on the v -> o head space. + R1 = r1_ext if external else _get_orthogonal_matrix(model.config.hidden_size, mode) + assert R1 is not None + rotations = {"R1": R1} + _rotate_input_cols(embed, R1) # writer: rows e <- e @ R1 + _rotate_input_cols(model.lm_head, R1) # reader: W <- W @ R1 (untied above) + for idx, layer in enumerate(layers): + attn, mlp = layer.self_attn, layer.mlp + if use_r2: + if external: + assert r2_ext is not None + R2 = r2_ext[idx] + else: + R2 = _get_orthogonal_matrix(head_dim, mode) + else: + R2 = None + _rotate_input_cols(attn.q_proj, R1) + _rotate_input_cols(attn.k_proj, R1) + _rotate_input_cols(attn.v_proj, R1) + _rotate_output_rows(attn.o_proj, R1) + _rotate_input_cols(mlp.gate_proj, R1) + _rotate_input_cols(mlp.up_proj, R1) + _rotate_output_rows(mlp.down_proj, R1) + # SpinQuant additionally folds the weight half of the ONLINE R4 activation Hadamard + # into down_proj; without the matching runtime activation transform that destroys + # the model, so it is deliberately skipped in this offline-only transform. + if use_r2: + _rotate_v_proj_r2(attn.v_proj, R2, head_dim) + _rotate_o_proj_r2(attn.o_proj, R2, head_dim) + rotations[f"model.layers.{idx}.self_attn.R2"] = R2 + + # 5. Post-conditions: the invariants that make the transform an identity. + for idx, layer in enumerate(layers): + for norm_name, _ in spec["norm_edges"]: + assert torch.all(layer.get_submodule(norm_name).weight.data == 1), ( + f"layer {idx} {norm_name} not fused to ones" + ) + if spec["has_qk_norm"]: + assert torch.equal( + layer.self_attn.q_norm.weight.data, qk_norm_before[f"{idx}.q_norm"] + ), f"q_norm[{idx}] changed" + assert torch.equal( + layer.self_attn.k_norm.weight.data, qk_norm_before[f"{idx}.k_norm"] + ), f"k_norm[{idx}] changed" + assert torch.all(decoder.norm.weight.data == 1), "final norm not fused to ones" + for n, p in model.named_parameters(): + assert tuple(p.shape) == shapes_before[n], f"shape of {n} changed" + + return {k: v.cpu() for k, v in rotations.items()} + + +def fold_seam_diags(model: nn.Module, seam_diags, smax: float = 256.0) -> dict: + """Bake learned per-input-channel seam scales into ``model`` in place (fp64 math). + + The transform-QAT counterpart of :meth:`fold_rotations` for the diagonal half of the + learned reparametrization (``RotationSet.seam_diags``): exactly the two + ROTATION-SURVIVING SmoothQuant-style prefold seams, as exact per-seam + functional identities — + + - **down seam** (``s_down [intermediate_size]``): ``up_proj`` rows ``/= s_down`` + (SwiGLU is elementwise: ``silu(g) * (u/s) == (silu(g)*u)/s``), ``down_proj`` + cols ``*= s_down``. + - **o seam** (``s_o [n_kv_heads*head_dim]``, GQA-exact — a per-channel scale + commutes with the per-head convex attention mix, shared across the q-heads of one + kv head): ``v_proj`` rows ``/= s_o``, ``o_proj`` cols ``*= s_o`` expanded per + q-head group. + + Row-scaled linears (up/v) have any bias divided too; col-scaled ones (down/o) are + input-side, bias untouched. All math is float64, cast back to each weight's dtype — + the model function is unchanged up to that round-trip. Composes with + :meth:`fold_rotations` in EITHER order (each fold is an identity); applying + ``fold_seam_diags`` FIRST reproduces the learner's effective-weight assembly + (prefold-inside, rotation-outside) entry-for-entry. + + Args: + model: Registered HF causal LM (``LlamaForCausalLM``/``Qwen3ForCausalLM``), + modified in place. + seam_diags: Mapping layer index -> ``{"down": s_down, "o": s_o}`` with strictly + positive scale vectors (the :attr:`RotationSet.seam_diags` format; keys may + be int or int-like str). May cover a subset of layers — each layer's seams + are independent identities. + smax: fp16-safety ceiling: scales are clamped to ``[1e-4, smax]`` before folding + (default 256 — the fp16-endpoint-safe value from the + activation-underflow finding, vs. 1e4 for bf16-only paths). A clamp that + actually bites trades exactness for numeric safety and is reported in the + returned evidence. + + Returns: + Evidence dict: ``{"smax": ..., "layers": {i: {"down_s_max", "down_s_spread", + "o_s_max", "o_s_spread", "clamped"}}}``. + """ + arch = type(model).__name__ + if arch not in _ARCH_REGISTRY: + raise NotImplementedError( + f"fold_seam_diags: unsupported architecture {arch!r}; " + f"supported: {sorted(_ARCH_REGISTRY)}" + ) + head_dim = _ARCH_REGISTRY[arch]["head_dim"](model.config) + layers = model.model.layers + + by_idx: dict[int, dict] = {} + for k, pair in seam_diags.items(): + idx = int(k) + if not 0 <= idx < len(layers): + raise ValueError(f"seam_diags layer index {idx} out of range 0..{len(layers) - 1}") + if set(pair) != {"down", "o"}: + raise ValueError(f"seam_diags[{k!r}]: expected keys {{'down', 'o'}}, got {set(pair)}") + by_idx[idx] = pair + + def _prep(vec, dim: int, name: str) -> tuple[torch.Tensor, torch.Tensor]: + """Validate one scale vector; return ``(clamped fp64 CPU scales, raw scales)``. + + Telemetry reports the RAW stats so a biting clamp is visible. + """ + s = torch.as_tensor(vec).detach().to(torch.float64).cpu().flatten() + if s.numel() != dim: + raise ValueError(f"{name}: expected {dim} scales, got {s.numel()}") + if not bool((s > 0).all()): + raise ValueError(f"{name}: scales must be strictly positive") + return s.clamp(1e-4, smax), s + + evidence: dict = {"smax": smax, "layers": {}} + for idx in sorted(by_idx): + layer = layers[idx] + up, down = layer.mlp.up_proj.weight, layer.mlp.down_proj.weight + v, o = layer.self_attn.v_proj.weight, layer.self_attn.o_proj.weight + + s_d, s_d_raw = _prep(by_idx[idx]["down"], down.shape[1], f"seam_diags[{idx}]['down']") + assert up.shape[0] == s_d.numel(), ( + f"layer {idx}: up_proj out {up.shape[0]} != down_proj in {s_d.numel()}" + ) + s_o, s_o_raw = _prep(by_idx[idx]["o"], v.shape[0], f"seam_diags[{idx}]['o']") + assert s_o.numel() % head_dim == 0, ( + f"layer {idx}: o-seam scale dim {s_o.numel()} not divisible by head_dim {head_dim}" + ) + assert o.shape[1] % s_o.numel() == 0, ( + f"layer {idx}: o_proj in {o.shape[1]} not a multiple of v_proj out {s_o.numel()}" + ) + n_kv = s_o.numel() // head_dim + group = o.shape[1] // s_o.numel() # q-heads per kv head (GQA) + + # down seam: up rows / s_d (bias too — row scaling), down cols * s_d. + s_dd = s_d.to(up.device) + up.data = (up.data.to(torch.float64) / s_dd[:, None]).to(up.dtype) + if layer.mlp.up_proj.bias is not None: + b = layer.mlp.up_proj.bias + b.data = (b.data.to(torch.float64) / s_dd).to(b.dtype) + down.data = (down.data.to(torch.float64) * s_dd[None, :]).to(down.dtype) + + # o seam: v rows / s_o (bias too), o cols * s_o expanded per q-head group. + s_oo = s_o.to(v.device) + v.data = (v.data.to(torch.float64) / s_oo[:, None]).to(v.dtype) + if layer.self_attn.v_proj.bias is not None: + b = layer.self_attn.v_proj.bias + b.data = (b.data.to(torch.float64) / s_oo).to(b.dtype) + s_full = s_oo.reshape(n_kv, 1, head_dim).expand(n_kv, group, head_dim).reshape(-1) + o.data = (o.data.to(torch.float64) * s_full[None, :].to(o.device)).to(o.dtype) + + evidence["layers"][idx] = { + "down_s_max": s_d_raw.max().item(), + "down_s_spread": (s_d_raw.max() / s_d_raw.min()).item(), + "o_s_max": s_o_raw.max().item(), + "o_s_spread": (s_o_raw.max() / s_o_raw.min()).item(), + "clamped": bool((s_d != s_d_raw).any()) or bool((s_o != s_o_raw).any()), + } + return evidence diff --git a/modelopt/torch/quantization/rotation/learn.py b/modelopt/torch/quantization/rotation/learn.py new file mode 100644 index 00000000000..8449d289630 --- /dev/null +++ b/modelopt/torch/quantization/rotation/learn.py @@ -0,0 +1,1115 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Learned SpinQuant rotations (R1 + per-layer R2) via Cayley SGD on the Stiefel manifold. + +:meth:`learn_rotations` optimizes the same offline rotations that :meth:`fold_rotations` +folds — a global residual-stream rotation R1 ``[hidden, hidden]`` and one per-layer +head-space rotation R2 ``[head_dim, head_dim]`` on the v_proj -> o_proj path — by minimizing +the next-token cross-entropy of the *fake-quantized* rotated model on calibration text +(SpinQuant, https://arxiv.org/abs/2405.16406). The rotation parameters live on the Stiefel +manifold and are updated with a Cayley-transform SGD (the SGDG optimizer of Li et al., +ported self-contained in :mod:`.sgdg`), so every iterate stays orthogonal to numerical tolerance and +the result folds through :meth:`fold_rotations`'s validated path via its ``R1=`` / ``R2=`` +arguments. + +Design notes, hyperparameter lineage (official SpinQuant vs. our internal reference trainer +vs. this module) and the objective-config table live in README.md next to this file. Architecture +knowledge (norm-fusion edges, reader/writer orientation, Qwen3 q/k_norm exclusion, +head_dim resolution, tied-embedding handling) is REUSED from ``fold.py`` — it is defined +exactly once, in the fold module's ``_ARCH_REGISTRY``. +""" + +import math +import random +import time +from collections.abc import Iterable, Mapping +from dataclasses import dataclass, field +from typing import Any + +import torch +import torch.nn as nn + +from .fold import _ARCH_REGISTRY, _fuse_norm_into_linears, _get_orthogonal_matrix +from .sgdg import SGDG + +__all__ = [ + "INT8_DEFAULT_OBJECTIVE", + "SEAM_DIAG_LR", + "SGDG", + "W4A4_G128_OBJECTIVE", + "W16A4_ASYM_R4G_OBJECTIVE", + "QuantObjective", + "RotationSet", + "learn_rotations", +] + +# Deployability tolerance for rotation matrices entering a fold. Trained (fp32 Cayley) +# iterates drift off the manifold, and the max-entry residual is BASIS-DEPENDENT: +# ``R^T R - I`` and ``R R^T - I`` share eigenvalues but not entries — measured on real +# 150-step runs the R^T R form sits at ~5e-5 while the R R^T form (the one the fold +# orientation actually consumes: reader/writer seams compose to ``x R1 R1^T W^T``) is +# 10-20x larger (~1e-3). learn_rotations therefore applies a FINAL RETRACTION (polar +# projection to the nearest orthogonal matrix, ~1e-14 residual both forms) before +# returning, so its outputs pass this gate with orders-of-magnitude headroom; the 1e-4 +# tolerance exists for hand-supplied matrices. Raw legacy R.bins (no retraction) need +# ``RotationSet.load(path, orthogonalize=True)``. +LEARNED_ORTHO_TOL = 1e-4 + +#: Peak Adam learning rate for the seam-diagonal parameters (``log s``) when +#: ``QuantObjective.learn_seam_diag`` is on. The diagonals live in a SEPARATE plain-Adam +#: group (never in the SGDG stiefel group — they are not on the manifold) and follow the +#: same cosine schedule as the rotation lr. +SEAM_DIAG_LR = 1e-2 + +#: Key under which :attr:`RotationSet.seam_diags` is stored inside a saved rotation +#: file. Absent from old-format (pure-R.bin) files, so those load with +#: ``seam_diags=None`` unchanged; a flat R.bin consumer must pop/ignore this key. +_SEAM_DIAGS_KEY = "__seam_diags__" + +# The per-layer linears whose (rotated) weights the objective fake-quantizes and whose +# inputs the activation quantizer sees. Matches ModelOpt PTQ coverage for decoder LMs: +# embeddings and lm_head are never quantized (INT8_DEFAULT_CFG disables ``*lm_head*``). +_ATTN_PROJS = ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.o_proj") +_MLP_PROJS = ("mlp.gate_proj", "mlp.up_proj", "mlp.down_proj") + + +# -------------------------------------------------------------------------------------- +# Fake-quant objective configuration (pluggable; STE numerics follow ModelOpt convention) +# -------------------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class QuantObjective: + """Fake-quantization spec for the training objective. + + Symmetric integer quant-dequant with ModelOpt max-calibration numerics: for ``b`` bits + the scale is ``amax / (2**(b-1) - 1)`` (clamped to 1e-12), values are rounded + half-to-even and clamped to ``[-2**(b-1), 2**(b-1) - 1]`` (e.g. int4 -> ``amax/7``, + ``[-8, 7]``; int8 -> ``amax/127``, ``[-128, 127]``). Gradients pass straight through. + + Args: + name: Tag recorded in :attr:`RotationSet.meta`. + w_bits: Weight bit-width, or None to disable weight fake-quant. + w_group: Weight quant group size along in_features (per-group), or None for + per-output-channel scales. + a_bits: Activation bit-width, or None to disable activation fake-quant. + a_mode: ``"per_token_dynamic"`` (scale from each token's amax, recomputed every + forward — SpinQuant/QuaRot W4A4 regime) or ``"per_tensor_static"`` (one scalar + scale per linear — the ModelOpt INT8_DEFAULT_CFG ``algorithm="max"`` axis). + a_static_scope: How the per-tensor-static scale surrogate is formed during + training (deployment always recalibrates on the FINAL folded model): + ``"batch"`` (default) — fresh amax per calibration batch, i.e. the scale + tracks the current rotation exactly as post-hoc max calibration would; the + stationary surrogate for the deployed endpoint. ``"run"`` — monotone running + max over all batches seen (literal max-calibration semantics over the calib + stream) — measured to be NON-STATIONARY under a moving rotation: past + rotations' outliers keep the scale inflated, later steps train against grids + deployment never uses, and the loss can drift UP (observed on Qwen3-1.7B: + 3.75 -> 4.87 over 150 steps). Kept for ablation. + learn_seam_diag: OSTQuant-style transform-QAT — additionally learn per-input- + channel diagonal scales at the two ROTATION-SURVIVING seams (down_proj input + and o_proj input; at every norm-fed seam a folded diagonal is cancelled by + norm fusion, so those seams have no surviving degree of freedom). The scales + are parametrized as ``log s`` (init 0 = identity, positivity for free) and + applied in the effective-weight assembly exactly like the SmoothQuant + prefold: ``up_proj`` rows ``/ s_down`` + ``down_proj`` cols ``* s_down``, + and ``v_proj`` rows ``/ s_o`` (KV dim) + ``o_proj`` cols ``* s_o`` expanded + per q-head group (GQA-exact) — a functional identity for ANY positive + diagonal, trained jointly with R1/R2 through the same STE objective but in a + separate plain-Adam group (:data:`SEAM_DIAG_LR`, same cosine schedule). + Default False: bitwise-identical behavior to the rotation-only trainer. + a_asym: Per-token dynamic ASYMMETRIC min-max affine activation fake-quant + instead of the symmetric default — the official SpinQuant activation + quantizer (paper A.4: "asymmetric quantization outperforms symmetric ... + no clipping"). Numerics match the official ``ActQuantizer`` (``sym=False``, + ``clip_ratio=1``) exactly: zero-inclusive token range, all-zero-token + fallback to ``[-1, 1]``, ``scale = (max - min)/(2**b - 1)``, + ``zp = round(-min/scale)``. Only implemented for + ``a_mode="per_token_dynamic"``. + r4_in_graph: Put the online R4 down_proj Hadamard into the TRAINING graph only + (input hook ``x @ H`` before activation fake-quant + effective down_proj + weight columns ``@ H`` — a functional-identity pair). The official trainer + does this unconditionally, even for no-had deployment + (``train_utils/main.py``); our ablation measured that training WITHOUT it + (deployment-faithful) makes the no-had result worse. The deployed model + folds R1/R2 only — no online op survives in the returned + :class:`RotationSet` or the fold. Power-of-2 seam dimension only. + """ + + name: str + w_bits: int | None = 4 + w_group: int | None = 128 + a_bits: int | None = None + a_mode: str = "per_token_dynamic" + a_static_scope: str = "batch" + learn_seam_diag: bool = False + a_asym: bool = False + r4_in_graph: bool = False + + def __post_init__(self): + # Bit-widths below 2 have no representable positive level: the symmetric scale is + # amax / (2**(b-1) - 1) = amax / 0 = inf, and the dequant multiply then yields + # 0 * inf = NaN for every value, silently NaN-ing the whole objective. + for field_name in ("w_bits", "a_bits"): + bits = getattr(self, field_name) + if bits is not None and bits < 2: + raise ValueError( + f"{field_name}={bits}: bit-width must be >= 2 (b=1 gives a zero " + "quantization range and an all-NaN fake-quant)" + ) + if self.a_mode not in ("per_token_dynamic", "per_tensor_static"): + raise ValueError(f"unknown a_mode: {self.a_mode!r}") + if self.a_static_scope not in ("batch", "run"): + raise ValueError(f"unknown a_static_scope: {self.a_static_scope!r}") + if self.a_asym and self.a_bits is None: + raise ValueError("a_asym=True requires a_bits") + if self.a_asym and self.a_mode != "per_token_dynamic": + raise ValueError("a_asym is only implemented for a_mode='per_token_dynamic'") + + +#: SpinQuant-paper-style W4A4 (per-group-128 sym weights + per-token dynamic sym int4 +#: activations). Comparison point for the paper's W4A4 rows. +W4A4_G128_OBJECTIVE = QuantObjective( + name="w4a4_g128", w_bits=4, w_group=128, a_bits=4, a_mode="per_token_dynamic" +) + +#: The axes of ModelOpt's INT8_DEFAULT_CFG: per-output-channel sym int8 weights + +#: per-tensor STATIC sym int8 activations ("max" calibration), lm_head excluded. This is +#: the deployment target where random rotations barely help — the static per-tensor +#: activation scale is the collapse axis the learned rotation trains against. +#: ``a_static_scope="batch"``: the scale tracks the current rotation (see QuantObjective). +INT8_DEFAULT_OBJECTIVE = QuantObjective( + name="int8_default", w_bits=8, w_group=None, a_bits=8, a_mode="per_tensor_static" +) + +#: The official trainer's objective for GPTQ-deployed rows ("Cayley on 16-4-KV", paper +#: Table 3): weights stay 16-bit in the loss, A4 per-token dynamic ASYM min-max (paper +#: A.4), and the online R4 down_proj Hadamard lives in the TRAINING graph only — the +#: official code keeps it there even when deploying no-had, and our ablation measured that +#: removing it (deployment-faithful training) makes the no-had result WORSE. The deployed +#: model still folds R1/R2 only. +W16A4_ASYM_R4G_OBJECTIVE = QuantObjective( + name="w16a4_asym_r4g", + w_bits=None, + w_group=None, + a_bits=4, + a_mode="per_token_dynamic", + a_asym=True, + r4_in_graph=True, +) + + +class _FakeQuantSTE(torch.autograd.Function): + """Symmetric quant-dequant with a precomputed scale; straight-through backward.""" + + @staticmethod + def forward(ctx, x, scale, qneg, qpos): + return (torch.round(x / scale).clamp_(qneg, qpos)) * scale + + @staticmethod + def backward(ctx, grad_out): + return grad_out, None, None, None + + +def _fq_weight(w: torch.Tensor, cfg: QuantObjective) -> torch.Tensor: + """Fake-quantize a ``[out, in]`` weight per ``cfg`` (per-group or per-out-channel).""" + qpos = float(2 ** (cfg.w_bits - 1) - 1) + qneg = -qpos - 1.0 + if cfg.w_group is None: + s = (w.abs().amax(dim=1, keepdim=True) / qpos).clamp_min(1e-12) + return _FakeQuantSTE.apply(w, s, qneg, qpos) + out_f, in_f = w.shape + g = cfg.w_group + wg = w.reshape(out_f, in_f // g, g) + s = (wg.abs().amax(dim=-1, keepdim=True) / qpos).clamp_min(1e-12) + return _FakeQuantSTE.apply(wg, s, qneg, qpos).reshape(out_f, in_f) + + +def _fq_act(x: torch.Tensor, s: torch.Tensor, bits: int) -> torch.Tensor: + qpos = float(2 ** (bits - 1) - 1) + q = (torch.round(x / s).clamp_(-qpos - 1.0, qpos)) * s + return x + (q - x).detach() # STE + + +def _fq_act_asym(x: torch.Tensor, bits: int) -> torch.Tensor: + """Per-token dynamic ASYM min-max affine fake-quant (STE backward). + + Bit-for-bit the official ``ActQuantizer`` (``sym=False``, ``clip_ratio=1``) recipe: + zero-inclusive token range, all-zero tokens fall back to ``[-1, 1]``, + ``scale = (xmax - xmin)/(2**bits - 1)``, ``zp = round(-xmin/scale)``, + ``q = clamp(round(x/scale) + zp, 0, 2**bits - 1)``, dequant ``scale * (q - zp)``. + """ + qmax = float(2**bits - 1) + xd = x.detach() + xmin = xd.amin(dim=-1, keepdim=True).clamp_max_(0.0) + xmax = xd.amax(dim=-1, keepdim=True).clamp_min_(0.0) + degen = (xmin == 0) & (xmax == 0) + xmin = torch.where(degen, torch.full_like(xmin, -1.0), xmin) + xmax = torch.where(degen, torch.full_like(xmax, 1.0), xmax) + s = (xmax - xmin) / qmax + zp = torch.round(-xmin / s) + q = (torch.round(xd / s) + zp).clamp_(0.0, qmax) + dq = (q - zp) * s + return x + (dq - x).detach() # STE + + +def _walsh_hadamard(n: int, device=None) -> torch.Tensor: + """Normalized Sylvester Walsh-Hadamard matrix ``[n, n]`` (fp32): symmetric, ``H Hᵀ = I``. + + The fixed (not random-signed) Hadamard the official online-R4 op applies + (``matmul_hadU``). Power-of-2 sizes only — the R4 seam dimension is + ``config.intermediate_size`` (Llama-3.2-1B: 8192 = 2^13); non-power-of-2 seams + (e.g. Qwen3's 6144 = 3·2048) need the had-K Kronecker composition, not implemented. + """ + if n <= 0 or (n & (n - 1)) != 0: + raise NotImplementedError( + f"r4_in_graph: seam dimension {n} is not a power of 2; the had-K Kronecker " + "composition is not implemented" + ) + H = torch.ones(1, 1, dtype=torch.float32, device=device) + while H.shape[0] < n: + H = torch.cat([torch.cat([H, H], dim=1), torch.cat([H, -H], dim=1)], dim=0) + return H / math.sqrt(n) + + +class _ActQuantHooks: + """Forward-pre-hooks applying STE activation fake-quant on the target linears. + + ``per_token_dynamic``: scale = token amax / qpos, recomputed each forward. + ``per_tensor_static``: one scalar scale per linear. Scope ``"batch"``: fresh amax per + batch — the scale tracks the current rotation, the stationary surrogate for post-hoc + max calibration of the folded model. Scope ``"run"``: monotone running max over every + batch seen (literal calib-stream max semantics; non-stationary under a moving + rotation — see QuantObjective). ``static_amax`` telemetry records the observed + per-linear maximum either way. + + ``r4_had`` (``QuantObjective.r4_in_graph``): the mlp.down_proj hook additionally + applies ``x @ H`` BEFORE its activation fake-quant — the online half of the + training-graph R4 whose weight half lives in the assembly. Pass H pre-cast to the + model dtype. With ``a_bits=None`` (W16-act-only variants keep a_bits set; this is + the r4-only corner) the hooks attach to down_proj alone and just rotate. + """ + + def __init__(self, cfg: QuantObjective, r4_had: torch.Tensor | None = None): + self.cfg = cfg + self.r4_had = r4_had + self.handles: list = [] + self.static_amax: dict[str, torch.Tensor] = {} + + def _hook(self, name: str): + cfg = self.cfg + r4h = self.r4_had if name.endswith("mlp.down_proj") else None + qpos = None if cfg.a_bits is None else float(2 ** (cfg.a_bits - 1) - 1) + + def hook(module, inputs): + x = inputs[0] + if r4h is not None: + x = x @ (r4h if r4h.dtype == x.dtype else r4h.to(x.dtype)) + if cfg.a_bits is None: + return (x, *inputs[1:]) + if cfg.a_mode == "per_token_dynamic": + if cfg.a_asym: + return (_fq_act_asym(x, cfg.a_bits), *inputs[1:]) + s = (x.detach().abs().amax(dim=-1, keepdim=True) / qpos).clamp_min(1e-12) + else: # per_tensor_static + batch_amax = x.detach().abs().amax() + prev = self.static_amax.get(name) + run_max = batch_amax if prev is None else torch.maximum(prev, batch_amax) + self.static_amax[name] = run_max # telemetry: observed max either way + amax = batch_amax if cfg.a_static_scope == "batch" else run_max + s = (amax / qpos).clamp_min(1e-12) + return (_fq_act(x, s, cfg.a_bits), *inputs[1:]) + + return hook + + def attach(self, model: nn.Module) -> int: + targets = ( + _ATTN_PROJS + _MLP_PROJS + if self.cfg.a_bits is not None + else ("mlp.down_proj",) # r4-only: nothing to quantize elsewhere + ) + for name, module in model.named_modules(): + if isinstance(module, nn.Linear) and ".layers." in name and name.endswith(targets): + self.handles.append(module.register_forward_pre_hook(self._hook(name))) + return len(self.handles) + + def remove(self): + for h in self.handles: + h.remove() + self.handles.clear() + + +# -------------------------------------------------------------------------------------- +# Final retraction — nearest orthogonal matrix (the SGDG optimizer itself lives in .sgdg) +# -------------------------------------------------------------------------------------- + + +def _polar_project(R: torch.Tensor) -> torch.Tensor: + """Nearest orthogonal matrix in Frobenius norm: ``R = U S V^T -> U V^T`` (float64). + + Used as the FINAL retraction on trained matrices: fp32 Cayley iterates accumulate + manifold drift whose max-entry residual is much larger in the ``R R^T`` form than the + ``R^T R`` form the trainer audits (measured 10-20x on 150-step runs), and the fold + orientation consumes ``R R^T``. The projection moves each entry by O(drift/2) — far + below one bf16 ulp for the measured ~1e-3 drift — and restores exact orthogonality + (~1e-14 both forms), making the subsequent checkpoint fold an exact identity again. + This is the same retraction semantics SGDG itself applies stochastically + (:meth:`_qr_retraction` with p = 1/101), applied deterministically once at the end. + """ + U, _, Vh = torch.linalg.svd(R.to(torch.float64)) + return U @ Vh + + +# -------------------------------------------------------------------------------------- +# RotationSet — the learned matrices, fold-ready +# -------------------------------------------------------------------------------------- + + +@dataclass +class RotationSet: + """Learned rotations in the fold/R.bin key convention, plus training telemetry. + + ``rotations`` maps ``"R1"`` and ``"model.layers.{i}.self_attn.R2"`` to float64 CPU + matrices — exactly the dict :meth:`fold_rotations` returns and the format SpinQuant's + optimized-rotation checkpoints (``R.bin``) use. Feed the matrices back through + ``fold_rotations(model, R1=rs.R1, R2=rs.R2)``. + + ``seam_diags`` (transform-QAT, ``QuantObjective.learn_seam_diag``) maps layer index + -> ``{"down": s_down [intermediate], "o": s_o [n_kv_heads*head_dim]}`` — the learned + per-input-channel seam SCALES (``exp`` of the trained log-parameters: strictly + positive; identity = ones) as float64 CPU vectors, or None when the diagonals were + not learned. Bake them into a fresh model with + :meth:`~modelopt.torch.quantization.rotation.fold_seam_diags` (plus + ``fold_rotations`` for R). ``save``/``load`` round-trip them; old-format R.bin files + load with ``seam_diags=None``. + """ + + rotations: dict[str, torch.Tensor] + history: list[dict] = field(default_factory=list) + meta: dict[str, Any] = field(default_factory=dict) + seam_diags: dict[int, dict[str, torch.Tensor]] | None = None + + def __post_init__(self): + if "R1" not in self.rotations: + raise ValueError("RotationSet requires an 'R1' entry") + # .clone() is load-bearing: the preceding conversions are no-ops for an already- + # float64 CPU input, so without it the stored rotations would alias the caller's + # buffers and could change after validation (and after save()). + self.rotations = { + k: torch.as_tensor(v).detach().to(torch.float64).cpu().clone() + for k, v in self.rotations.items() + } + if self.seam_diags is not None: + norm = {} + for k, pair in self.seam_diags.items(): + if set(pair) != {"down", "o"}: + raise ValueError( + f"seam_diags[{k!r}]: expected keys {{'down', 'o'}}, got {set(pair)}" + ) + if int(k) in norm: + raise ValueError(f"seam_diags names layer {int(k)} more than once") + norm[int(k)] = { + kk: torch.as_tensor(vv).detach().to(torch.float64).cpu().flatten().clone() + for kk, vv in pair.items() + } + for kk, vv in norm[int(k)].items(): + if not bool((vv > 0).all()): + raise ValueError( + f"seam_diags[{k!r}][{kk!r}]: scales must be strictly positive" + ) + self.seam_diags = norm + + @property + def R1(self) -> torch.Tensor: + """The global residual-stream rotation (the ``"R1"`` entry).""" + return self.rotations["R1"] + + @property + def R2(self) -> dict[str, torch.Tensor]: + """The per-layer R2 sub-dict (keys ``model.layers.{i}.self_attn.R2``).""" + return {k: v for k, v in self.rotations.items() if k != "R1"} + + def ortho_audit(self) -> dict[str, float]: + """Per rotation: ``max(max |R^T R - I|, max |R R^T - I|)`` in float64. + + Both forms are measured because their max-entry residuals differ for + near-orthogonal matrices (basis-dependent; measured 10-20x on raw trained R1), + and the fold orientation consumes the ``R R^T`` form. + """ + out = {} + for k, R in self.rotations.items(): + eye = torch.eye(R.shape[0], dtype=torch.float64) + out[k] = max( + (R.t() @ R - eye).abs().max().item(), + (R @ R.t() - eye).abs().max().item(), + ) + return out + + def save(self, path) -> None: + """Write the flat float64 rotation dict (R.bin-compatible; telemetry not saved). + + With learned ``seam_diags`` present, they ride along under the reserved + :data:`_SEAM_DIAGS_KEY` entry (new format); ``seam_diags=None`` writes exactly + the legacy flat dict, byte-format-identical to before. + """ + payload: dict[str, Any] = dict(self.rotations) + if self.seam_diags is not None: + payload[_SEAM_DIAGS_KEY] = { + int(k): {kk: vv.clone() for kk, vv in pair.items()} + for k, pair in self.seam_diags.items() + } + torch.save(payload, path) + + @classmethod + def load( + cls, path, ortho_tol: float = LEARNED_ORTHO_TOL, orthogonalize: bool = False + ) -> "RotationSet": + """Load a flat rotation dict; refuse matrices off the manifold beyond ortho_tol. + + ``orthogonalize=True`` applies the polar retraction to every matrix before the + gate — for raw legacy R.bins written without the final retraction (their + ``R R^T`` residual is typically ~1e-3 and would be refused otherwise). + + Backward compatible both ways: old-format files (pure rotation dict) load with + ``seam_diags=None``; new-format files carry the learned seam diagonals under + :data:`_SEAM_DIAGS_KEY` (never polar-projected — they are not rotations). + """ + raw = torch.load(path, map_location="cpu", weights_only=True) + seam_diags = raw.pop(_SEAM_DIAGS_KEY, None) + if orthogonalize: + raw = {k: _polar_project(torch.as_tensor(v)) for k, v in raw.items()} + rs = cls(rotations=raw, seam_diags=seam_diags) + for k, err in rs.ortho_audit().items(): + if err >= ortho_tol: + raise ValueError( + f"{path}: rotation {k!r} is not orthogonal " + f"(max ortho residual = {err:.3e} >= {ortho_tol}) — refusing to load" + " (raw trained R.bin? pass orthogonalize=True for the polar" + " retraction)" + ) + return rs + + +# -------------------------------------------------------------------------------------- +# Effective-weight assembly (differentiable mirror of fold.py's orientation table) +# -------------------------------------------------------------------------------------- + + +def _assemble_effective_weights( + base: dict[str, torch.Tensor], + R1: torch.Tensor, + R2s: list[torch.Tensor], + n_layers: int, + head_dim: int, + objective: QuantObjective | None, + out_dtype: torch.dtype, + seam_diag_params: list[dict[str, torch.Tensor]] | None = None, + r4_had: torch.Tensor | None = None, +) -> dict[str, torch.Tensor]: + """Build every rotated (and weight-fake-quantized) effective weight for one forward. + + Same orientation table as fold.py — readers ``W @ R1``, writers ``R1^T @ W``, + embed/lm_head included, v_proj per-KV-head rows ``R2^T @ W_h``, o_proj per-Q-head + columns ``@ R2`` — but computed out-of-place in R's dtype (fp32) with R1/R2 as graph + leaves, so ``loss.backward()`` reaches the rotation parameters. Because every + residual-stream seam is consistently rotated, the assembled model computes the + offline-rotated model's function for ANY orthogonal R1/R2 (see README.md). + + ``seam_diag_params`` (transform-QAT): optional per-layer ``{"down": log_s_down + [intermediate], "o": log_s_o [n_kv_heads*head_dim]}`` graph leaves. Applied with the + prefold structure BEFORE weight fake-quant, prefold-inside / rotation-outside + (i.e. exactly the composition ``t14_sq_prefold`` -> ``fold_rotations``): up_proj + rows ``/ s_down``; down_proj cols ``* s_down``; v_proj rows ``/ s_o`` before the + per-KV-head R2 row step; o_proj cols ``* s_o`` expanded per q-head group before the + per-Q-head R2 column step. A functional identity for any positive diagonal (and a + bitwise no-op path when None), so gradients on ``log s`` come only from the + quantization error, like the rotations'. + + ``r4_had`` (``QuantObjective.r4_in_graph``): the normalized Walsh-Hadamard for the + down_proj seam — the effective down_proj weight gets its columns rotated (``@ H``), + pairing with the input-side ``x @ H`` hook into a functional identity that only the + quantizers can see. Training-graph only: the fold consumes the returned R1/R2 and + never sees H. + """ + compute = R1.dtype + d = head_dim + + def fin(w): + if objective is not None and objective.w_bits is not None: + w = _fq_weight(w, objective) + return w.to(out_dtype) + + eff = {} + for name in ("model.embed_tokens.weight", "lm_head.weight"): # never fake-quantized + eff[name] = (base[name].to(compute) @ R1).to(out_dtype) + + for i in range(n_layers): + R2 = R2s[i] + pre = f"model.layers.{i}." + sp = None if seam_diag_params is None else seam_diag_params[i] + s_down = None if sp is None else torch.exp(sp["down"].to(compute)) + s_o = None if sp is None else torch.exp(sp["o"].to(compute)) + + for proj in ("self_attn.q_proj", "self_attn.k_proj", "mlp.gate_proj"): + n = pre + proj + ".weight" + eff[n] = fin(base[n].to(compute) @ R1) # readers + + n = pre + "mlp.up_proj.weight" # reader (R1); down-seam rows / s_down + a = base[n].to(compute) @ R1 + if s_down is not None: + a = a / s_down[:, None] + eff[n] = fin(a) + + n = pre + "mlp.down_proj.weight" # writer (R1); down-seam cols * s_down + w = R1.t() @ base[n].to(compute) + if s_down is not None: + w = w * s_down[None, :] + if r4_had is not None: # training-graph R4: cols @ H, inverse of the x @ H hook + w = w @ r4_had.to(compute) + eff[n] = fin(w) + + n = pre + "self_attn.v_proj.weight" # reader (R1) + o-seam rows / s_o (KV dim) + a = base[n].to(compute) @ R1 + o_f, i_f = a.shape + if s_o is not None: + a = a / s_o[:, None] # before R2: prefold-inside, rotation-outside + a = (a.t().reshape(i_f, o_f // d, d) @ R2).reshape(i_f, o_f).t().contiguous() + eff[n] = fin(a) + + n = pre + "self_attn.o_proj.weight" # writer (R1) + o-seam cols * s_o expanded + w = R1.t() @ base[n].to(compute) + o_f, i_f = w.shape + if s_o is not None: + n_kv = s_o.numel() // d + group = i_f // s_o.numel() # q-heads per KV head (GQA sharing) + s_full = s_o.reshape(n_kv, 1, d).expand(n_kv, group, d).reshape(-1) + w = w * s_full[None, :] # before R2: prefold-inside, rotation-outside + eff[n] = fin((w.reshape(o_f, i_f // d, d) @ R2).reshape(o_f, i_f)) + return eff + + +def _iter_batches(calib_loader: Iterable, steps: int): + """Yield ``steps`` batches, restarting the loader between epochs if re-iterable.""" + if steps <= 0: + return + n = 0 + it = iter(calib_loader) + while n < steps: + try: + batch = next(it) + except StopIteration: + it = iter(calib_loader) + try: + batch = next(it) + except StopIteration: + raise ValueError( + "calib_loader yielded no batches (a one-shot generator that is already " + "exhausted? pass a re-iterable loader, e.g. a list or DataLoader)" + ) from None + yield batch + n += 1 + + +def _batch_input_ids(batch) -> tuple[torch.Tensor, torch.Tensor | None]: + """Extract ``(input_ids, attention_mask)`` from a calibration batch. + + The mask is returned (not dropped) so padded batches neither attend over padding nor + contribute pad positions to the cross-entropy; ``None`` means "no padding in this + batch". + """ + mask = None + if isinstance(batch, torch.Tensor): + ids = batch + elif isinstance(batch, Mapping) or hasattr(batch, "keys"): + ids = batch["input_ids"] + mask = batch.get("attention_mask") + else: + raise TypeError( + f"unsupported calib batch type {type(batch).__name__}: pass input_ids tensors " + "[bs, seq] or dicts with an 'input_ids' key" + ) + if ids.dim() == 1: + ids = ids.unsqueeze(0) + assert ids.dim() == 2, f"input_ids must be [bs, seq], got shape {tuple(ids.shape)}" + if mask is not None: + if mask.dim() == 1: + mask = mask.unsqueeze(0) + if mask.shape != ids.shape: + raise ValueError( + f"attention_mask shape {tuple(mask.shape)} != input_ids shape {tuple(ids.shape)}" + ) + return ids, mask + + +# -------------------------------------------------------------------------------------- +# Public API +# -------------------------------------------------------------------------------------- + + +def learn_rotations( + model: nn.Module, + calib_loader: Iterable, + steps: int = 150, + lr: float = 1.5, + mode: str = "hadamard", + objective_cfg: QuantObjective | None = W4A4_G128_OBJECTIVE, + seed: int = 0, + init_rotations: Mapping[str, torch.Tensor] | None = None, + log_every: int = 10, + teacher: nn.Module | None = None, + kd_alpha: float = 0.5, + kd_temp: float = 2.0, +) -> RotationSet: + """Learn SpinQuant rotations R1 + per-layer R2 for ``model`` by Cayley SGD. + + The model's weights are FROZEN; only the rotations train. Each step assembles the + rotated effective weights out-of-place (same orientation table as + :meth:`fold_rotations`), applies the objective's weight fake-quant (STE), runs the + frozen model reparametrized with those weights on a calibration batch (activation + fake-quant applied by pre-hooks per the objective), and takes one SGDG step on the + next-token cross-entropy with cosine lr decay — the SpinQuant recipe. + + Side effects on ``model`` (identical to fold_rotations' pre-rotation steps, so a + subsequent ``fold_rotations(model, R1=..., R2=...)`` on the SAME object — or on a + fresh copy — is equivalent): tied embeddings are untied with a real clone + (``config.tie_word_embeddings`` set False), RMSNorm gains are fused into downstream + linears, and all parameters get ``requires_grad=False``. Weights are otherwise + unchanged (the rotated weights live only in the per-step reparametrization); + Qwen3 q/k_norm are bitwise untouched (asserted). Seeds the global torch CPU RNG and + Python's ``random`` (for SGDG's stochastic QR retraction). + + Args: + model: HuggingFace causal LM whose class is registered in fold.py's + ``_ARCH_REGISTRY`` (``LlamaForCausalLM``, ``Qwen3ForCausalLM``). Run on the + model's current device/dtype (CPU fp32 works; GPU bf16 recommended for real + models). + calib_loader: Re-iterable of calibration batches — ``input_ids`` tensors + ``[bs, seq]`` or dicts with an ``input_ids`` key. Cycled for ``steps`` steps. + steps: Number of Cayley-SGD steps (the paper reports quality saturating around + 100; 150 is our reference budget). ``steps=0`` returns + the untouched init (== ``fold_rotations(mode=mode, seed=seed)`` draws). + lr: Peak learning rate, cosine-decayed to 0 (official SpinQuant default 1.5). + mode: Rotation init family, ``"hadamard"`` (random-sign Hadamard) or ``"random"`` + (Haar QR) — same draw order as fold_rotations, so equal seeds give the same + init matrices. + objective_cfg: :class:`QuantObjective` for the fake-quant loss, or None to train + against the unquantized CE (near-zero gradient — the rotated model is exactly + equivalent; useful only as a null check). + seed: Seed for the init draws and the retraction RNG. + init_rotations: Optional warm start — a dict in the R.bin key convention + (overrides mode/seed draws; gated at :data:`LEARNED_ORTHO_TOL`). + log_every: Print a progress line every N steps (0 = silent). + teacher: Optional frozen reference model for a KD objective: loss becomes + ``(1-kd_alpha)*CE + kd_alpha*kd_temp^2*KL(student || teacher)`` with the + teacher's logits computed under ``no_grad`` on the same batch. The teacher + is NEVER reparametrized/fused/modified — pass a separate (typically bf16) + copy, not the model being trained. ``teacher=None`` (default) is the plain + CE objective, bitwise-identical to the plain-CE trainer. + kd_alpha: KD mixing weight in [0, 1] (only with ``teacher``). + kd_temp: KD softmax temperature > 0 (only with ``teacher``). + + Returns: + :class:`RotationSet` with float64 CPU matrices (audited orthonormal to + :data:`LEARNED_ORTHO_TOL`), per-step ``history`` and hyperparameter ``meta``. + With ``objective_cfg.learn_seam_diag=True`` (transform-QAT) it additionally + carries the jointly learned per-layer seam scales in + :attr:`RotationSet.seam_diags` — trained by a separate plain-Adam group + (:data:`SEAM_DIAG_LR`, same cosine schedule; never the SGDG stiefel group) and + bakeable into a fresh model via + :meth:`~modelopt.torch.quantization.rotation.fold_seam_diags`. + """ + from torch.nn.utils import stateless as _stateless + + arch = type(model).__name__ + if arch not in _ARCH_REGISTRY: + raise NotImplementedError( + f"learn_rotations: unsupported architecture {arch!r}; " + f"supported: {sorted(_ARCH_REGISTRY)}" + ) + spec: dict[str, Any] = _ARCH_REGISTRY[arch] + + if teacher is not None: + # A teacher that IS the student is reparametrized and hooked along with it, so its + # logits equal the student's and the KD term is identically zero — a plain-CE run + # scaled by (1 - kd_alpha) while meta reports KD as active. Pass a separate copy. + if teacher is model: + raise ValueError( + "teacher must be a distinct module from the model being trained (it is " + "reparametrized in-place during each step, so a self-teacher gives KL == 0); " + "pass e.g. copy.deepcopy(model) frozen in the reference dtype" + ) + if not 0.0 <= kd_alpha <= 1.0: + raise ValueError(f"kd_alpha must be in [0, 1], got {kd_alpha}") + if kd_temp <= 0.0: + raise ValueError(f"kd_temp must be > 0, got {kd_temp} (temperature divides logits)") + + decoder = model.model + layers = decoder.layers + embed = decoder.embed_tokens + n_layers = len(layers) + hidden = model.config.hidden_size + head_dim = spec["head_dim"](model.config) + device = embed.weight.device + model_dtype = embed.weight.dtype + t_all = time.time() + + # 1. Untie embeddings (real clone) — same as fold_rotations step 1. + if model.lm_head.weight.data_ptr() == embed.weight.data_ptr(): + model.lm_head.weight = nn.Parameter( + embed.weight.data.clone(), requires_grad=embed.weight.requires_grad + ) + model.config.tie_word_embeddings = False + + # Snapshots for post-conditions. + qk_norm_before = {} + if spec["has_qk_norm"]: + for idx, layer in enumerate(layers): + qk_norm_before[f"{idx}.q_norm"] = layer.self_attn.q_norm.weight.data.clone() + qk_norm_before[f"{idx}.k_norm"] = layer.self_attn.k_norm.weight.data.clone() + shapes_before = {n: tuple(p.shape) for n, p in model.named_parameters()} + + # 2. Fuse RMSNorm gains into downstream linears (fold_rotations step 3; fused norms + # become exactly ones, which is what lets R1 commute through RMSNorm). Fusing an + # already-fused model (all-ones gains) is a bitwise no-op, so this is idempotent. + for layer in layers: + for norm_name, linear_names in spec["norm_edges"]: + _fuse_norm_into_linears( + layer.get_submodule(norm_name), + [layer.get_submodule(n) for n in linear_names], + ) + _fuse_norm_into_linears(decoder.norm, [model.lm_head]) + + # 3. Freeze the model; rotations are the only trainable parameters. + model.eval() + model.requires_grad_(False) + + # 4. Rotation parameters: fp32 on-device, init = the SAME seeded draw order as + # fold_rotations (R1 first, then R2 by ascending layer), or a warm start. + torch.manual_seed(seed) + random.seed(seed) # SGDG's stochastic QR-retraction draw — reproducible trajectories + if init_rotations is None: + draws = {"R1": _get_orthogonal_matrix(hidden, mode)} + for i in range(n_layers): + draws[f"model.layers.{i}.self_attn.R2"] = _get_orthogonal_matrix(head_dim, mode) + else: + draws = { + k: torch.as_tensor(v).detach().to(torch.float64).cpu() + for k, v in init_rotations.items() + } + assert "R1" in draws and draws["R1"].shape == (hidden, hidden), ( + f"init_rotations: R1 missing or wrong shape (want {(hidden, hidden)})" + ) + for i in range(n_layers): + k = f"model.layers.{i}.self_attn.R2" + assert k in draws and draws[k].shape == (head_dim, head_dim), ( + f"init_rotations: missing/misshaped {k}" + ) + for k, R64 in draws.items(): + eye = torch.eye(R64.shape[0], dtype=torch.float64) + # BOTH residual forms: they share eigenvalues but not entries, and the fold + # consumes the R R^T form while the closing audit reports the max of the two. + # Gating on one alone lets a warm start pass here and fail the exit audit. + err = max( + (R64 @ R64.T - eye).abs().max().item(), + (R64.T @ R64 - eye).abs().max().item(), + ) + if err >= LEARNED_ORTHO_TOL: + raise ValueError( + f"init_rotations[{k!r}]: not orthogonal (max |R R^T - I|, |R^T R - I| = " + f"{err:.3e} >= {LEARNED_ORTHO_TOL}); pass " + "RotationSet.load(..., orthogonalize=True) to retract a legacy matrix" + ) + R1 = nn.Parameter(draws["R1"].to(device=device, dtype=torch.float32)) + R2s = [ + nn.Parameter(draws[f"model.layers.{i}.self_attn.R2"].to(device=device, dtype=torch.float32)) + for i in range(n_layers) + ] + + # 4b. Transform-QAT seam diagonals (QuantObjective.learn_seam_diag): per-layer + # log-scale vectors for the two rotation-surviving seams, init zeros (= identity + # scales, so the assembled model is function-preserving at init and NO extra RNG is + # consumed — the R draw/trajectory stream is unchanged either way). + learn_seam_diag = objective_cfg is not None and objective_cfg.learn_seam_diag + seam_diag_params: list[dict[str, nn.Parameter]] | None = None + if learn_seam_diag: + intermediate = model.config.intermediate_size + n_kv = model.config.num_key_value_heads + n_q = model.config.num_attention_heads + assert n_q % n_kv == 0, f"GQA group not integral: {n_q} q heads / {n_kv} kv heads" + seam_diag_params = [ + { + "down": nn.Parameter(torch.zeros(intermediate, dtype=torch.float32, device=device)), + "o": nn.Parameter(torch.zeros(n_kv * head_dim, dtype=torch.float32, device=device)), + } + for _ in range(n_layers) + ] + + # 5. Base-weight references (frozen, post-fusion) + objective preconditions. + sd = dict(model.named_parameters()) + base = { + "model.embed_tokens.weight": sd["model.embed_tokens.weight"].data, + "lm_head.weight": sd["lm_head.weight"].data, + } + for i in range(n_layers): + for proj in _ATTN_PROJS + _MLP_PROJS: + name = f"model.layers.{i}.{proj}.weight" + base[name] = sd[name].data + bias = f"model.layers.{i}.{proj}.bias" + assert bias not in sd, ( + f"{bias} exists — bias handling is not implemented in the assembly" + ) + if ( + objective_cfg is not None + and objective_cfg.w_bits is not None + and objective_cfg.w_group is not None + ): + assert base[name].shape[1] % objective_cfg.w_group == 0, ( + f"{name}: in_features {base[name].shape[1]} not divisible by " + f"w_group {objective_cfg.w_group}" + ) + + r4_had = None + if objective_cfg is not None and objective_cfg.r4_in_graph: + r4_had = _walsh_hadamard(model.config.intermediate_size, device=device) + assert base["model.layers.0.mlp.down_proj.weight"].shape[1] == r4_had.shape[0], ( + "down_proj in_features != config.intermediate_size" + ) + + # Built here but ATTACHED inside the try below: hooks registered before the guarded + # region would survive on the caller's model if any later setup step raised, silently + # fake-quantizing every subsequent forward. + hooks = None + expected_hooks = 0 + if objective_cfg is not None and (objective_cfg.a_bits is not None or r4_had is not None): + hooks = _ActQuantHooks( + objective_cfg, + r4_had=None if r4_had is None else r4_had.to(model_dtype), + ) + expected_hooks = ( + n_layers * len(_ATTN_PROJS + _MLP_PROJS) + if objective_cfg.a_bits is not None + else n_layers + ) + + if learn_seam_diag: + assert base["model.layers.0.mlp.up_proj.weight"].shape[0] == intermediate, ( + "up_proj out_features != config.intermediate_size" + ) + assert base["model.layers.0.self_attn.v_proj.weight"].shape[0] == n_kv * head_dim, ( + "v_proj out_features != num_key_value_heads * head_dim" + ) + assert base["model.layers.0.self_attn.o_proj.weight"].shape[1] == n_q * head_dim, ( + "o_proj in_features != num_attention_heads * head_dim" + ) + + # The rotations train on the Stiefel manifold (SGDG); the seam diagonals are + # UNCONSTRAINED log-parameters and get their own plain-Adam optimizer (never the + # stiefel group) at SEAM_DIAG_LR, sharing the cosine schedule below. + opt = SGDG([R1, *R2s], lr=lr, momentum=0.0, stiefel=True) + opt_diag = None + if learn_seam_diag: + assert seam_diag_params is not None + opt_diag = torch.optim.Adam( + [p for sp in seam_diag_params for p in (sp["down"], sp["o"])], + lr=SEAM_DIAG_LR, + ) + history: list[dict] = [] + + def _r1_ortho() -> float: + Rd = R1.detach().to(torch.float64) + eye = torch.eye(Rd.shape[0], dtype=torch.float64, device=Rd.device) + return (Rd.t() @ Rd - eye).abs().max().item() + + # 6. Training loop. + try: + if hooks is not None: + n_hooked = hooks.attach(model) + assert n_hooked == expected_hooks, ( + f"activation hooks attached to {n_hooked} linears, expected {expected_hooks}" + ) + for step, batch in enumerate(_iter_batches(calib_loader, steps)): + cos_t = 0.5 * (1.0 + math.cos(math.pi * step / max(steps, 1))) + lr_t = lr * cos_t + for gp in opt.param_groups: + gp["lr"] = lr_t + if opt_diag is not None: # same cosine schedule, SEAM_DIAG_LR peak + for gp in opt_diag.param_groups: + gp["lr"] = SEAM_DIAG_LR * cos_t + ids, attn_mask = _batch_input_ids(batch) + ids = ids.to(device) + if attn_mask is not None: + attn_mask = attn_mask.to(device) + # Padding must not enter the loss: HF ignores label index -100. + labels = ids.masked_fill(attn_mask == 0, -100) + else: + labels = ids + t0 = time.time() + eff = _assemble_effective_weights( + base, + R1, + R2s, + n_layers, + head_dim, + objective_cfg, + model_dtype, + seam_diag_params=seam_diag_params, + r4_had=r4_had, + ) + # Reparametrize for forward AND backward: with activation checkpointing the + # recompute happens during backward and must still see the effective weights + # (torch.func.functional_call would restore the originals first). + with _stateless._reparametrize_module(model, eff): + out = model(input_ids=ids, attention_mask=attn_mask, labels=labels, use_cache=False) + loss = out.loss + if teacher is not None: # KD objective; teacher never touched + with torch.no_grad(): + tlogits = teacher( + input_ids=ids, attention_mask=attn_mask, use_cache=False + ).logits + T = kd_temp + # Flatten to [tokens, vocab] so "batchmean" divides by the TOKEN count: + # on [bs, seq, vocab] it divides by bs alone, making the KD term seq_len + # times the per-token KL and the CE:KD mix sequence-length-dependent. + # Padding positions (label -100) are excluded from the KD term too. + keep = (labels != -100).reshape(-1) + s_logits = out.logits.reshape(-1, out.logits.shape[-1])[keep] + t_logits = tlogits.reshape(-1, tlogits.shape[-1])[keep] + kd = torch.nn.functional.kl_div( + torch.nn.functional.log_softmax(s_logits / T, dim=-1), + torch.nn.functional.softmax(t_logits / T, dim=-1), + reduction="batchmean", + ) * (T * T) + loss = (1.0 - kd_alpha) * loss + kd_alpha * kd + opt.zero_grad(set_to_none=True) + if opt_diag is not None: + opt_diag.zero_grad(set_to_none=True) + loss.backward() + del eff + if step == 0: # gradients must reach every R (and every diag) and nothing else + assert R1.grad is not None and all(r.grad is not None for r in R2s), ( + "no gradient reached the rotation parameters" + ) + if seam_diag_params is not None: + assert all( + sp[k].grad is not None for sp in seam_diag_params for k in ("down", "o") + ), "no gradient reached the seam-diagonal parameters" + assert all(p.grad is None for p in model.parameters()), ( + "a frozen model weight received a gradient" + ) + opt.step() + if opt_diag is not None: + opt_diag.step() + rec = { + "step": step, + "lr": round(lr_t, 6), + "loss": round(loss.item(), 6), + "r1_ortho": _r1_ortho(), + "dt_s": round(time.time() - t0, 3), + } + history.append(rec) + if log_every and (step % log_every == 0 or step == steps - 1): + print( + f"[learn_rotations] step {step:4d}/{steps} lr={lr_t:.4f} " + f"loss={rec['loss']:.4f} r1_ortho={rec['r1_ortho']:.2e}", + flush=True, + ) + finally: + if hooks is not None: + hooks.remove() + + # 7. Post-conditions: the model is functionally the original (untied + fused only). + if spec["has_qk_norm"]: + for idx, layer in enumerate(layers): + assert torch.equal( + layer.self_attn.q_norm.weight.data, qk_norm_before[f"{idx}.q_norm"] + ), f"q_norm[{idx}] changed" + assert torch.equal( + layer.self_attn.k_norm.weight.data, qk_norm_before[f"{idx}.k_norm"] + ), f"k_norm[{idx}] changed" + for n, p in model.named_parameters(): + assert tuple(p.shape) == shapes_before[n], f"shape of {n} changed" + + retraction_log: dict[str, dict[str, float]] = {} + if history: + # Final retraction: project each trained matrix to the nearest orthogonal one + # (see _polar_project). Raw fp32 Cayley drift is recorded per matrix (both + # residual forms + max entry moved) before being closed to ~1e-14. + rotations = {} + raw64 = {"R1": R1.detach().to(torch.float64).cpu()} + for i, r2 in enumerate(R2s): + raw64[f"model.layers.{i}.self_attn.R2"] = r2.detach().to(torch.float64).cpu() + for k, R in raw64.items(): + eye = torch.eye(R.shape[0], dtype=torch.float64) + proj = _polar_project(R) + retraction_log[k] = { + "raw_rtr": (R.t() @ R - eye).abs().max().item(), + "raw_rrt": (R @ R.t() - eye).abs().max().item(), + "delta_max": (proj - R).abs().max().item(), + } + rotations[k] = proj + else: + # steps=0: hand back the exact float64 init draws — bitwise identical to + # fold_rotations(mode=mode, seed=seed)'s matrices (same RNG, same draw order). + # No retraction: fresh draws are orthonormal to ~1e-15 by construction. + rotations = dict(draws) + + # Transform-QAT: export the learned seam scales s = exp(log_s) (fp64 CPU; ones when + # steps=0 — the identity init). Positivity is structural (exp), so these always pass + # the RotationSet gate. + seam_diags = None + if seam_diag_params is not None: + seam_diags = { + i: { + "down": torch.exp(sp["down"].detach().to(torch.float64)).cpu(), + "o": torch.exp(sp["o"].detach().to(torch.float64)).cpu(), + } + for i, sp in enumerate(seam_diag_params) + } + + meta = { + "arch": arch, + "hidden": hidden, + "head_dim": head_dim, + "n_layers": n_layers, + "steps": steps, + "lr": lr, + "mode": mode, + "seed": seed, + "objective": None if objective_cfg is None else objective_cfg.__dict__.copy(), + "kd": None if teacher is None else {"alpha": kd_alpha, "temp": kd_temp}, + "warm_start": init_rotations is not None, + "final_loss": history[-1]["loss"] if history else None, + "final_retraction": retraction_log, + "elapsed_s": round(time.time() - t_all, 1), + } + if hooks is not None and hooks.static_amax: + meta["static_act_amax"] = {k: float(v) for k, v in hooks.static_amax.items()} + if seam_diags is not None: + all_s = torch.cat([v for pair in seam_diags.values() for v in pair.values()]) + meta["seam_diag"] = { + "lr": SEAM_DIAG_LR, + "s_min": all_s.min().item(), + "s_max": all_s.max().item(), + } + + rs = RotationSet(rotations=rotations, history=history, meta=meta, seam_diags=seam_diags) + audit = rs.ortho_audit() + worst = max(audit, key=lambda k: audit[k]) + assert audit[worst] < LEARNED_ORTHO_TOL, ( + f"ORTHO AUDIT FAILED after training: max |R^T R - I| = {audit[worst]:.3e} " + f">= {LEARNED_ORTHO_TOL} ({worst}) — rotations are not deployable" + ) + return rs diff --git a/modelopt/torch/quantization/rotation/sgdg.py b/modelopt/torch/quantization/rotation/sgdg.py new file mode 100644 index 00000000000..1ea5c659c94 --- /dev/null +++ b/modelopt/torch/quantization/rotation/sgdg.py @@ -0,0 +1,219 @@ +# Adapted from https://github.com/JunLi-Galios/Optimization-on-Stiefel-Manifold-via-Cayley-Transform/blob/c5ab4e8/stiefel_optimizer.py +# (the SGDG "Cayley SGD" Stiefel-manifold optimizer of Li et al., ICLR 2020), with the +# repository's helper functions (gutils.py / utils.py) inlined and minor API and +# robustness modifications documented in the port note below. The same code is vendored +# as train_utils/optimizer.py in Meta's SpinQuant repository; this port was verified +# bitwise against that copy on the stiefel branch at momentum 0.0 and 0.9. +# Copyright (c) 2020 Jun Li +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 AND MIT +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SGDG — Cayley SGD on the Stiefel manifold (self-contained port, see header for lineage). + +The Stiefel-branch MATH is unchanged from the original — every rotation parameter used by +:mod:`.learn` is square fp32 with ``stiefel=True``, so that branch always runs. Textual +deltas vs. the original: + +- stiefel branch: three unused local reads (``weight_decay``/``dampening``/``nesterov``) + dropped; +- non-stiefel fallback: deprecated ``Tensor.add_(Number, Tensor)`` overloads use the + modern ``add_(t, alpha=...)`` signature, and the original's bare try/except around the + weight-decay add (which silently skipped weight decay when stale locals leaked from a + prior stiefel iteration) is replaced by fresh group reads. + +KNOWN ORIGINAL QUIRK, faithfully reproduced: momentum is INERT in the stiefel branch. +``V = momentum * V - g.t()`` rebinds ``V`` to a temp and ``V.copy_(V_new)`` writes into +that temp — ``state["momentum_buffer"]`` stays zeros forever (dead store; proven: momentum +0.9 vs 0.0 give bitwise-identical trajectories). Kept for trajectory parity with the +original code and our reference runs. + +NOTE: the stiefel branch draws ``random.randint`` from Python's GLOBAL random module +(original behavior, kept): with p = 1/101 per parameter per step the iterate is +re-projected by QR retraction. Seed ``random`` for reproducible trajectories +(:meth:`.learn.learn_rotations` does this from its ``seed`` argument). +""" + +import random + +import torch +from torch.optim.optimizer import Optimizer + +__all__ = ["SGDG"] + + +def _unit(v, dim: int = 1, eps: float = 1e-8): + vnorm = _norm(v, dim) + return v / vnorm.add(eps), vnorm + + +def _norm(v, dim: int = 1): + assert len(v.size()) == 2 + return v.norm(p=2, dim=dim, keepdim=True) + + +def _matrix_norm_one(W): + return torch.abs(W).sum(dim=0).max() + + +def _cayley_loop(X, W, tan_vec, t): + """Fixed-point iteration for the Cayley transform: Y = X + t*W*(X+Y)/2, 5 iterations.""" + Y = X + t * tan_vec + for _ in range(5): + Y = X + t * torch.matmul(W, 0.5 * (X + Y)) + return Y.t() + + +def _qr_retraction(tan_vec): # tan_vec: p-by-n, p <= n + [p, n] = tan_vec.size() + tan_vec.t_() + q, r = torch.linalg.qr(tan_vec) + d = torch.diag(r, 0) + ph = d.sign() + q *= ph.expand_as(q) + q.t_() + return q + + +_EPSILON = 1e-8 + + +class SGDG(Optimizer): + """SGD-G: SGD on the Stiefel manifold via the Cayley transform (see module docstring). + + With ``stiefel=True`` and a square parameter, each step (i) row-normalizes the iterate, + (ii) builds the skew-symmetric tangent ``W = W_hat - W_hat^T`` from the Riemannian + gradient projection, (iii) moves along the Cayley curve with step ``min(lr, 1 / + ||W||_1)`` via 5 fixed-point iterations — an orthogonality-preserving update to + first order. With p = 1/101 per step the iterate is exactly re-orthogonalized by QR. + """ + + def __init__( + self, + params, + lr: float, + momentum: float = 0, + dampening: float = 0, + weight_decay: float = 0, + nesterov: bool = False, + stiefel: bool = False, + omega: float = 0, + grad_clip=None, + ) -> None: + """Set up parameter groups with the original SGDG defaults (see class docstring).""" + defaults = { + "lr": lr, + "momentum": momentum, + "dampening": dampening, + "weight_decay": weight_decay, + "nesterov": nesterov, + "stiefel": stiefel, + "omega": 0, + "grad_clip": grad_clip, + } + if nesterov and (momentum <= 0 or dampening != 0): + raise ValueError("Nesterov momentum requires a momentum and zero dampening") + super().__init__(params, defaults) + + def __setstate__(self, state) -> None: + """Restore optimizer state, defaulting ``nesterov`` for groups pickled without it.""" + super().__setstate__(state) + for group in self.param_groups: + group.setdefault("nesterov", False) + + @torch.no_grad() + def step(self, closure=None): + """Perform one optimization step (Cayley-curve update on the stiefel branch).""" + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + momentum = group["momentum"] + stiefel = group["stiefel"] + + for p in group["params"]: + if p.grad is None: + continue + + unity, _ = _unit(p.data.view(p.size()[0], -1)) + if stiefel and unity.size()[0] <= unity.size()[1]: + rand_num = random.randint(1, 101) + if rand_num == 1: + unity = _qr_retraction(unity) + + g = p.grad.data.view(p.size()[0], -1) + + lr = group["lr"] + + param_state = self.state[p] + if "momentum_buffer" not in param_state: + param_state["momentum_buffer"] = torch.zeros(g.t().size(), device=p.device) + + V = param_state["momentum_buffer"] + V = momentum * V - g.t() + MX = torch.mm(V, unity) + XMX = torch.mm(unity, MX) + XXMX = torch.mm(unity.t(), XMX) + W_hat = MX - 0.5 * XXMX + W = W_hat - W_hat.t() + t = 0.5 * 2 / (_matrix_norm_one(W) + _EPSILON) + alpha = min(t, lr) + + p_new = _cayley_loop(unity.t(), W, V, alpha) + V_new = torch.mm(W, unity.t()) # n-by-p + p.data.copy_(p_new.view(p.size())) + # Original dead store, reproduced verbatim: V was rebound to a temp + # above, so this never reaches state["momentum_buffer"] — momentum is + # inert in the stiefel branch (see module docstring). + V.copy_(V_new) + + else: + weight_decay = group["weight_decay"] + dampening = group["dampening"] + nesterov = group["nesterov"] + d_p = p.grad.data + if weight_decay != 0: + d_p.add_(p.data, alpha=weight_decay) + if momentum != 0: + param_state = self.state[p] + if "momentum_buffer" not in param_state: + buf = param_state["momentum_buffer"] = d_p.clone() + else: + buf = param_state["momentum_buffer"] + buf.mul_(momentum).add_(d_p, alpha=1 - dampening) + d_p = d_p.add(buf, alpha=momentum) if nesterov else buf + p.data.add_(d_p, alpha=-group["lr"]) + + return loss diff --git a/pyproject.toml b/pyproject.toml index e090fb62f86..837375d51fe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -222,6 +222,8 @@ extend-ignore = [ "*/_[a-zA-Z]*" = ["D"] # Private packages (_abc/*.py) or modules (_xyz.py) "*.ipynb" = ["D", "E501"] # Ignore missing docstrings or line length for Jupyter notebooks "modelopt/torch/kernels/*" = ["N803", "N806", "E731"] # triton style +"modelopt/torch/quantization/rotation/*" = ["N802", "N803", "N806"] # SpinQuant math notation (R1/R2 rotations, Hadamard H, Cayley tangent W) +"tests/unit/torch/quantization/test_rotation_*.py" = ["N802", "N803", "N806", "PERF203"] # same math notation; PERF203: the standalone __main__ drivers run each test in a try/except loop by design "modelopt/torch/puzzletron/*" = [ "C4", "D", diff --git a/tests/unit/torch/quantization/test_rotation_contracts.py b/tests/unit/torch/quantization/test_rotation_contracts.py new file mode 100644 index 00000000000..01f3932c5c2 --- /dev/null +++ b/tests/unit/torch/quantization/test_rotation_contracts.py @@ -0,0 +1,382 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Input-contract and state-hygiene regression tests for the rotation module. + +Each test pins one contract whose violation is SILENT (wrong numbers, a mutated caller +model, a dropped user input) rather than loud, i.e. exactly the class of defect the rest +of the suite's equivalence gates cannot catch. +""" + +import os + +os.environ["CUDA_VISIBLE_DEVICES"] = "" # CPU-only unit tests: never claim a GPU + +import sys +import traceback + +import pytest +import torch +from transformers import LlamaConfig, LlamaForCausalLM + +from modelopt.torch.quantization.rotation import QuantObjective, fold_rotations, learn_rotations + +VOCAB, HIDDEN, HEAD_DIM = 128, 64, 32 + +TINY_W4A4 = QuantObjective( + name="tiny_w4a4", w_bits=4, w_group=16, a_bits=4, a_mode="per_token_dynamic" +) + + +def _tiny_llama(): + """Two-layer Llama with non-unit RMSNorm gains (so norm fusion is observable).""" + torch.manual_seed(1234) + cfg = LlamaConfig( + vocab_size=VOCAB, + hidden_size=HIDDEN, + intermediate_size=2 * HIDDEN, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=HEAD_DIM, + max_position_embeddings=256, + tie_word_embeddings=False, + attn_implementation="eager", + ) + model = LlamaForCausalLM(cfg).eval() + for module in model.modules(): + if type(module).__name__.endswith("RMSNorm"): + module.weight.data = 1.0 + 0.1 * torch.randn_like(module.weight.data) + return model + + +def _batches(n=1, bs=2, seq=16, seed=7): + torch.manual_seed(seed) + return [torch.randint(0, VOCAB, (bs, seq)) for _ in range(n)] + + +def _orthogonal(n, seed): + torch.manual_seed(seed) + q, _ = torch.linalg.qr(torch.randn(n, n, dtype=torch.float64)) + return q.contiguous() + + +# -------------------------------------------------------------------------------------- +# 1. KD objective contracts +# -------------------------------------------------------------------------------------- + + +def test_kd_term_is_per_token_not_per_sequence(): + """The KD term must not scale with calibration sequence length. + + ``kl_div(reduction="batchmean")`` on unflattened [bs, seq, vocab] logits divides by bs + alone, making the KD term seq_len times the per-token KL — so the documented + ``(1-alpha)*CE + alpha*T^2*KL`` mix would silently change meaning when the caller + changes the calibration sequence length. ``kd_alpha=1.0`` makes the recorded loss + exactly the KD term, so the two lengths are directly comparable. + """ + + def kd_only_loss(seq): + torch.manual_seed(21) + teacher = _tiny_llama() + for p in teacher.parameters(): + p.data.add_(0.05 * torch.randn_like(p.data)) + teacher.eval().requires_grad_(False) + torch.manual_seed(4) + ids = torch.randint(0, VOCAB, (2, seq)) + rs = learn_rotations( + _tiny_llama(), + [ids], + steps=1, + lr=0.0, + objective_cfg=None, + seed=3, + log_every=0, + teacher=teacher, + kd_alpha=1.0, + kd_temp=2.0, + ) + return rs.history[0]["loss"] + + small, big = kd_only_loss(8), kd_only_loss(32) + assert big < 2.0 * small, ( + f"KD term scales with sequence length: seq=32 -> {big:.4f} vs seq=8 -> {small:.4f} " + "(a per-token mean keeps these comparable)" + ) + + +def test_self_teacher_rejected(): + """teacher=model must raise, not silently degenerate to a scaled plain-CE run. + + The teacher forward runs inside the student's reparametrization with the student's + activation-quant hooks attached, so a self-teacher yields identical logits and KL == 0 + while ``meta["kd"]`` still advertises KD as active. + """ + model = _tiny_llama() + with pytest.raises(ValueError, match="distinct module"): + learn_rotations( + model, + _batches(), + steps=1, + lr=0.0, + objective_cfg=TINY_W4A4, + seed=3, + log_every=0, + teacher=model, + kd_alpha=0.5, + ) + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"kd_temp": 0.0}, "kd_temp"), + ({"kd_alpha": 1.5}, "kd_alpha"), + ({"kd_alpha": -0.1}, "kd_alpha"), + ], +) +def test_kd_hyperparameters_validated(kwargs, match): + """kd_temp <= 0 (division by zero -> NaN much later) and kd_alpha outside [0, 1] + (which silently flips the CE term's sign) must be rejected at the call, not surface as + a NaN in the closing SVD.""" + with pytest.raises(ValueError, match=match): + learn_rotations( + _tiny_llama(), + _batches(), + steps=1, + lr=0.0, + objective_cfg=None, + seed=3, + log_every=0, + teacher=_tiny_llama(), + **kwargs, + ) + + +# -------------------------------------------------------------------------------------- +# 2. Calibration-batch contracts +# -------------------------------------------------------------------------------------- + + +def test_attention_mask_is_honored_in_dict_batches(): + """Padding must affect neither attention nor the loss. + + The documented dict-batch form carries ``attention_mask``; dropping it makes the CE a + function of how much padding the tokenizer happened to add, so identical real tokens + with different padding would train against different objectives. + """ + torch.manual_seed(11) + real = torch.randint(1, VOCAB, (2, 8)) + + def padded(n_pad): + ids = torch.cat([real, torch.zeros(2, n_pad, dtype=real.dtype)], dim=1) + mask = torch.cat( + [torch.ones(2, 8, dtype=torch.long), torch.zeros(2, n_pad, dtype=torch.long)], dim=1 + ) + return [{"input_ids": ids, "attention_mask": mask}] + + losses = [ + learn_rotations( + _tiny_llama(), padded(n_pad), steps=1, lr=0.0, objective_cfg=None, seed=3, log_every=0 + ).history[0]["loss"] + for n_pad in (2, 16) + ] + rel = abs(losses[0] - losses[1]) / max(abs(losses[0]), 1e-9) + assert rel < 0.02, ( + f"attention_mask ignored: 2 vs 16 pad tokens on identical real tokens give " + f"{losses[0]:.4f} vs {losses[1]:.4f} ({rel:.1%} apart)" + ) + + +def test_dict_batch_without_mask_matches_tensor_batch(): + """The documented dict form must be equivalent to the plain-tensor form when there is + no padding (the Mapping branch of the batch extractor is otherwise untested).""" + ids = _batches()[0] + as_tensor = learn_rotations( + _tiny_llama(), [ids], steps=1, lr=0.0, objective_cfg=None, seed=3, log_every=0 + ).history[0]["loss"] + as_dict = learn_rotations( + _tiny_llama(), + [{"input_ids": ids}], + steps=1, + lr=0.0, + objective_cfg=None, + seed=3, + log_every=0, + ).history[0]["loss"] + assert as_tensor == pytest.approx(as_dict, rel=1e-9) + + +def test_unsupported_batch_type_raises(): + """A batch that is neither a tensor nor a mapping must name the accepted forms.""" + with pytest.raises(TypeError, match="input_ids"): + learn_rotations( + _tiny_llama(), + [["not", "ids"]], + steps=1, + lr=0.0, + objective_cfg=None, + seed=3, + log_every=0, + ) + + +def test_empty_calib_loader_raises(): + """An exhausted/empty loader must be reported, not loop forever or divide by zero.""" + with pytest.raises(ValueError, match="no batches"): + learn_rotations(_tiny_llama(), [], steps=1, lr=0.0, objective_cfg=None, seed=3, log_every=0) + + +# -------------------------------------------------------------------------------------- +# 3. QuantObjective validation +# -------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize("field", ["w_bits", "a_bits"]) +def test_bit_width_below_two_rejected(field): + """b=1 gives scale = amax/(2^0 - 1) = amax/0 = inf and then 0*inf = NaN for every + value, so the whole objective silently becomes NaN; it must be rejected instead.""" + with pytest.raises(ValueError, match=field): + QuantObjective(name="bad", **{field: 1}) + + +def test_a_mode_validated_even_without_a_bits(): + """An unknown a_mode must be caught regardless of a_bits, so a typo cannot sit + dormant in a config until activations are switched on.""" + with pytest.raises(ValueError, match="a_mode"): + QuantObjective(name="bad", a_mode="per_tensor_dynamic") + + +# -------------------------------------------------------------------------------------- +# 4. External-input hygiene: no aliasing, no silent key collapse +# -------------------------------------------------------------------------------------- + + +def test_returned_rotations_do_not_alias_caller_buffers(): + """The returned provenance record must be a copy. + + Every conversion in the accept path (as_tensor/detach/to/cpu) is a no-op for an + already-float64 CPU input, so without an explicit clone the audited matrix would keep + changing whenever the caller reuses its buffer — after the orthogonality gate passed. + """ + r1 = _orthogonal(HIDDEN, seed=9) + rots = fold_rotations(_tiny_llama(), R1=r1, use_r2=False) + returned = rots["R1"] + assert returned.data_ptr() != r1.data_ptr() + before = returned.clone() + r1.mul_(2.0) + assert torch.equal(returned, before) + + +def test_duplicate_r2_layer_keys_rejected(): + """An int key and the R.bin-convention key naming the SAME layer must raise. + + Both normalize to one index, so last-writer-wins silently discards one user-supplied + rotation while the per-layer completeness check still passes — a wrong-rotation + checkpoint with no diagnostic. + """ + model = _tiny_llama() + r2 = { + 0: _orthogonal(HEAD_DIM, seed=1), + "model.layers.0.self_attn.R2": _orthogonal(HEAD_DIM, seed=2), + 1: _orthogonal(HEAD_DIM, seed=3), + } + with pytest.raises(ValueError, match="more than once"): + fold_rotations(model, R1=_orthogonal(HIDDEN, seed=4), R2=r2) + + +def test_warm_start_gate_checks_both_residual_forms(): + """A warm start must be rejected up front by the same criterion the exit audit uses. + + ``R^T R - I`` and ``R R^T - I`` share eigenvalues but not entries, so a matrix can + pass a one-sided gate and then fail the closing two-sided audit *after* the whole run + has been paid for. + """ + n = HIDDEN + torch.manual_seed(13) + q, _ = torch.linalg.qr(torch.randn(n, n, dtype=torch.float64)) + skew = torch.zeros(n, n, dtype=torch.float64) + skew[0, 1] = 1.2e-4 + r = q @ (torch.eye(n, dtype=torch.float64) + skew) + eye = torch.eye(n, dtype=torch.float64) + assert (r @ r.T - eye).abs().max().item() < 1e-4 <= (r.T @ r - eye).abs().max().item() + + model = _tiny_llama() + init = {"R1": r} + for i in range(model.config.num_hidden_layers): + init[f"model.layers.{i}.self_attn.R2"] = _orthogonal(HEAD_DIM, seed=100 + i) + with pytest.raises(ValueError, match="not orthogonal"): + learn_rotations( + model, + _batches(), + steps=0, + lr=0.0, + objective_cfg=None, + seed=3, + log_every=0, + init_rotations=init, + ) + + +# -------------------------------------------------------------------------------------- +# 5. State hygiene on the failure path +# -------------------------------------------------------------------------------------- + + +def test_failed_call_leaves_no_activation_hooks(): + """A setup failure must not leave fake-quant hooks on the caller's model. + + Hooks registered before the guarded region survive the exception, so every later + forward of that model silently quantizes activations and a retry double-quantizes. + """ + model = _tiny_llama() + ids = _batches()[0] + ref = model(input_ids=ids).logits.clone() + + bad = QuantObjective( + name="bad_seam", + w_bits=4, + w_group=16, + a_bits=4, + a_mode="per_token_dynamic", + learn_seam_diag=True, + ) + model.config.intermediate_size += 8 # provoke a seam-shape assert during setup + with pytest.raises(Exception): + learn_rotations(model, [ids], steps=1, lr=0.0, objective_cfg=bad, seed=3, log_every=0) + + after = model(input_ids=ids).logits + assert torch.allclose(ref, after, atol=1e-5), ( + "activation fake-quant hooks survived a failed learn_rotations call: " + f"max logit delta {(ref - after).abs().max().item():.3e}" + ) + + +if __name__ == "__main__": + tests = [(n, f) for n, f in sorted(globals().items()) if n.startswith("test_") and callable(f)] + failed = [] + for name, fn in tests: + try: + if hasattr(fn, "pytestmark"): # parametrized: run via pytest instead + print(f"SKIP {name} (parametrized)", flush=True) + continue + fn() + print(f"PASS {name}", flush=True) + except Exception: + failed.append(name) + print(f"FAIL {name}", flush=True) + traceback.print_exc() + print(f"\n{len(tests) - len(failed)} checked; FAILED: {failed}" if failed else "\nall passed") + sys.exit(1 if failed else 0) diff --git a/tests/unit/torch/quantization/test_rotation_ext_fold.py b/tests/unit/torch/quantization/test_rotation_ext_fold.py new file mode 100644 index 00000000000..a687b86edc4 --- /dev/null +++ b/tests/unit/torch/quantization/test_rotation_ext_fold.py @@ -0,0 +1,619 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Extensive fold property tests for +modelopt.torch.quantization.rotation.fold_rotations. + +Coverage beyond test_rotation_fold.py: + 1. an architecture sweep over {Llama, Qwen3} x {GQA 1,2,4} x {head_dim ==/!= hidden/heads} + x {tied, untied} x 2 fold seeds x {hadamard (incl. the Paley had-12 branch), random} + with the full identity-invariant bundle checked per cell; + 2. external-path == seed-path bitwise on every sweep config; + 3. idempotent norm fusion (direct re-fusion AND a full identity-matrix refold); + 4. R.bin round-trip of fold-returned dicts through RotationSet save/load, including the + off-manifold refusal and the orthogonalize=True polar retraction; + 5. use_r2=False on every config, and R2-as-list == R2-as-str-dict == R2-as-int-dict; + 6. an orientation oracle on a hand-built 1-layer model: reader W@R1 / writer R1^T W / + v-rows R2 / o-cols R2 verified against independent fp64 manual matmuls of the tiny + network (weight-level and seam-level), plus the RMSNorm rotation-invariance identity. + +Plain test_* functions with asserts: collectable by pytest, and also runnable without it +via ``python test_rotation_ext_fold.py`` (the __main__ driver runs every test function and +exits nonzero on any failure). CPU-only, tiny models. +""" + +import os + +os.environ["CUDA_VISIBLE_DEVICES"] = "" # CPU-only: never touch a GPU (a GPU chain runs) + +import sys +import tempfile +import traceback +from dataclasses import dataclass + +import pytest +import torch +from transformers import LlamaConfig, LlamaForCausalLM, Qwen3Config, Qwen3ForCausalLM + +from modelopt.torch.quantization.rotation import RotationSet, fold_rotations +from modelopt.torch.quantization.rotation.fold import _fuse_norm_into_linears + +VOCAB = 96 + +# FP-equivalence tolerance: same budget as test_rotation_fold.py (fold math is float64; +# error enters only via the fp64 -> fp32 cast per rewritten weight; measured max |delta +# logit| ~3e-7 on these sizes, so 1e-4 keeps >100x headroom while still catching any real +# transform bug, which perturbs logits by O(1)). +ATOL_FP32 = 1e-4 + + +# -------------------------------------------------------------------------------------- +# Config sweep: deterministic specs spanning the requested axes +# -------------------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class _Spec: + name: str + arch: str # "llama" | "qwen3" + hidden: int + n_layers: int + n_heads: int + n_kv: int + head_dim: int + tie: bool + mode: str # rotation mode for the seed path + + @property + def gqa_ratio(self) -> int: + return self.n_heads // self.n_kv + + @property + def head_dim_coincident(self) -> bool: + return self.head_dim == self.hidden // self.n_heads + + +# Spanning set (~8 tiny configs). head_dim is DELIBERATELY decoupled from hidden/heads on +# most cells (like Qwen3-0.6B: 128 vs 1024/16 = 64) so a head_dim-resolution regression to +# the wrong fallback formula cannot pass silently. hidden 48 / head_dim 24 exercise the +# Paley had-12 branch of the hadamard generator; hidden 36 / head_dim 12 have NO 2^k*K +# Hadamard decomposition, forcing (and covering) the mode="random" path. +_CONFIGS = ( + _Spec("llama_g1_hdEQ", "llama", 64, 2, 4, 4, 16, tie=False, mode="hadamard"), + _Spec("llama_g2_hdNE_tied", "llama", 64, 2, 4, 2, 32, tie=True, mode="hadamard"), + _Spec("llama_g4_hdEQ_rand", "llama", 32, 3, 4, 1, 8, tie=False, mode="random"), + _Spec("llama_g2_paley_tied", "llama", 48, 1, 4, 2, 24, tie=True, mode="hadamard"), + _Spec("qwen3_g1_hdNE", "qwen3", 64, 2, 4, 4, 32, tie=False, mode="hadamard"), + _Spec("qwen3_g4_hdEQ_tied", "qwen3", 64, 2, 8, 2, 8, tie=True, mode="hadamard"), + _Spec("qwen3_g1_odd_rand_tied", "qwen3", 36, 1, 2, 2, 12, tie=True, mode="random"), + _Spec("qwen3_g2_paley", "qwen3", 48, 2, 4, 2, 16, tie=False, mode="hadamard"), +) + +_FOLD_SEEDS = (0, 1) + + +def _randomize_rmsnorm_gains(model): + """Set every RMSNorm gain to a random non-one value: fresh HF models initialize norm + weights to ones, which would make the norm-fusion math and the fused-to-ones / + q_norm-untouched checks vacuous.""" + for module in model.modules(): + if type(module).__name__.endswith("RMSNorm"): + module.weight.data = 1.0 + 0.1 * torch.randn_like(module.weight.data) + + +def _build(spec: _Spec): + """Deterministic build: equal specs give bitwise-identical models (fixed init seed).""" + torch.manual_seed(1234) + kwargs = { + "vocab_size": VOCAB, + "hidden_size": spec.hidden, + "intermediate_size": 2 * spec.hidden, + "num_hidden_layers": spec.n_layers, + "num_attention_heads": spec.n_heads, + "num_key_value_heads": spec.n_kv, + "head_dim": spec.head_dim, + "max_position_embeddings": 128, + "tie_word_embeddings": spec.tie, + "attn_implementation": "eager", + } + if spec.arch == "llama": + model = LlamaForCausalLM(LlamaConfig(**kwargs)).eval() + else: + model = Qwen3ForCausalLM(Qwen3Config(**kwargs)).eval() + _randomize_rmsnorm_gains(model) + return model + + +def _logits(model, vocab=VOCAB): + torch.manual_seed(99) + ids = torch.randint(0, vocab, (2, 8)) + with torch.no_grad(): + return model(ids).logits + + +def _r2_keys(n_layers): + return {f"model.layers.{i}.self_attn.R2" for i in range(n_layers)} + + +def test_spec_table_spans_axes(): + """Meta-test: the sweep table really spans every requested axis, so the sweep tests + below cannot silently lose coverage if the table is edited.""" + assert len(_CONFIGS) == 8 + assert len({s.name for s in _CONFIGS}) == 8 + for arch in ("llama", "qwen3"): + sub = [s for s in _CONFIGS if s.arch == arch] + assert {s.gqa_ratio for s in sub} == {1, 2, 4}, f"{arch}: GQA ratios not spanned" + assert {s.head_dim_coincident for s in sub} == {True, False}, ( + f"{arch}: need both head_dim == and != hidden/heads" + ) + assert {s.tie for s in sub} == {True, False}, f"{arch}: need tied and untied" + assert {s.mode for s in _CONFIGS} == {"hadamard", "random"} + assert len(_FOLD_SEEDS) == 2 + # the Paley (48 = 2^2*12) and no-Hadamard-decomposition (36) cells are present + assert any(s.hidden == 48 and s.mode == "hadamard" for s in _CONFIGS) + assert any(s.hidden == 36 and s.mode == "random" for s in _CONFIGS) + + +# -------------------------------------------------------------------------------------- +# 1. Randomized architecture sweep: identity invariants on every config x seed +# -------------------------------------------------------------------------------------- + + +def test_arch_sweep_fold_identity(): + """For every spanning config x 2 fold seeds: fold_rotations (seed path) keeps fp32 + logits equal within ATOL_FP32; every pre-existing parameter shape is unchanged (the + only allowed new parameter is the untied lm_head.weight, embed-shaped); Qwen3 + q/k_norm are bitwise untouched; fused norms are exactly ones; the returned dict has + the R.bin key set with correctly-sized orthonormal fp64 CPU matrices (R2 sized + head_dim x head_dim — catches head_dim-resolution regressions per config).""" + for spec in _CONFIGS: + for fold_seed in _FOLD_SEEDS: + tag = f"{spec.name} seed={fold_seed}" + model = _build(spec) + shapes_before = {n: tuple(p.shape) for n, p in model.named_parameters()} + embed_before = model.model.embed_tokens.weight.data.clone() + qk_before = { + n: p.data.clone() + for n, p in model.named_parameters() + if "q_norm" in n or "k_norm" in n + } + if spec.arch == "qwen3": + assert len(qk_before) == 2 * spec.n_layers, tag + assert all(not torch.all(v == 1) for v in qk_before.values()), ( + f"{tag}: q/k_norm gains not randomized — untouched check would be vacuous" + ) + else: + assert not qk_before, tag + + before = _logits(model) + rots = fold_rotations(model, mode=spec.mode, seed=fold_seed, use_r2=True) + after = _logits(model) + + # (a) fp32 functional identity + max_diff = (after - before).abs().max().item() + assert torch.allclose(after, before, rtol=0, atol=ATOL_FP32), ( + f"{tag}: max |delta logit| = {max_diff:.3e} > {ATOL_FP32}" + ) + # ... and the fold really rewrote weights (guards against a vacuous no-op) + assert not torch.equal(model.model.embed_tokens.weight.data, embed_before), ( + f"{tag}: embed_tokens unchanged — fold was a no-op" + ) + + # (b) parameter shapes unchanged; only allowed new param: untied lm_head + params_after = dict(model.named_parameters()) + for n, shp in shapes_before.items(): + assert n in params_after, f"{tag}: parameter {n} disappeared" + assert tuple(params_after[n].shape) == shp, f"{tag}: shape of {n} changed" + new = set(params_after) - set(shapes_before) + if spec.tie: + assert new == {"lm_head.weight"}, f"{tag}: unexpected new params {new}" + assert tuple(params_after["lm_head.weight"].shape) == tuple( + model.model.embed_tokens.weight.shape + ), f"{tag}: untied lm_head has wrong shape" + else: + assert not new, f"{tag}: unexpected new params {new}" + + # (c) Qwen3 q/k_norm bitwise untouched + for n, v in qk_before.items(): + assert torch.equal(params_after[n].data, v), f"{tag}: {n} changed" + + # (d) fused norms exactly ones + for li, layer in enumerate(model.model.layers): + assert torch.all(layer.input_layernorm.weight == 1), ( + f"{tag}: layer {li} input_layernorm not ones" + ) + assert torch.all(layer.post_attention_layernorm.weight == 1), ( + f"{tag}: layer {li} post_attention_layernorm not ones" + ) + assert torch.all(model.model.norm.weight == 1), f"{tag}: final norm not ones" + + # (e) returned rotations: key set, dtype/device, per-config sizes, orthonormal + assert set(rots) == {"R1"} | _r2_keys(spec.n_layers), f"{tag}: bad key set" + for k, R in rots.items(): + size = spec.hidden if k == "R1" else spec.head_dim + assert R.dtype == torch.float64 and R.device.type == "cpu", f"{tag}: {k}" + assert R.shape == (size, size), ( + f"{tag}: {k} shape {tuple(R.shape)} != {(size, size)}" + ) + err = (R @ R.T - torch.eye(size, dtype=torch.float64)).abs().max().item() + assert err < 1e-10, f"{tag}: {k} max |R R^T - I| = {err:.3e}" + + +# -------------------------------------------------------------------------------------- +# 2. External path == seed path bitwise on every config +# -------------------------------------------------------------------------------------- + + +def test_external_path_bitwise_matches_seed_path(): + """For every sweep config: feeding the seed path's returned matrices back through + R1=/R2= on an identically-constructed model reproduces every parameter bitwise, and + the external fold returns the same matrices bitwise.""" + for spec in _CONFIGS: + m_seed = _build(spec) + rots = fold_rotations(m_seed, mode=spec.mode, seed=0, use_r2=True) + m_ext = _build(spec) # identical construction seed -> identical weights + returned = fold_rotations( + m_ext, R1=rots["R1"], R2={k: v for k, v in rots.items() if k != "R1"} + ) + ps = dict(m_seed.named_parameters()) + pe = dict(m_ext.named_parameters()) + assert set(ps) == set(pe), f"{spec.name}: parameter sets differ" + for n, p in pe.items(): + assert torch.equal(p.data, ps[n].data), f"{spec.name}: {n} differs" + assert set(returned) == set(rots), f"{spec.name}: returned key set differs" + for k in rots: + assert torch.equal(returned[k], rots[k]), f"{spec.name}: returned {k} differs" + + +# -------------------------------------------------------------------------------------- +# 3. Idempotent norm fusion +# -------------------------------------------------------------------------------------- + + +def test_idempotent_norm_fusion_refold(): + """On an already-folded model (all norm gains exactly ones): (a) re-running the norm + fusion directly is a bitwise no-op on every weight; (b) a full identity refold + (external R1 = I, R2 = I) — which re-runs untie + fusion + all rotation applications — + is a bitwise no-op on every parameter (fp64 multiply by 1.0 / matmul by I is exact).""" + for spec in (_CONFIGS[1], _CONFIGS[4]): # llama tied + qwen3 untied + model = _build(spec) + fold_rotations(model, mode=spec.mode, seed=0, use_r2=True) + snap = {n: p.data.clone() for n, p in model.named_parameters()} + + # (a) direct re-fusion of every already-fused edge + for layer in model.model.layers: + attn, mlp = layer.self_attn, layer.mlp + _fuse_norm_into_linears(layer.input_layernorm, [attn.q_proj, attn.k_proj, attn.v_proj]) + _fuse_norm_into_linears(layer.post_attention_layernorm, [mlp.gate_proj, mlp.up_proj]) + _fuse_norm_into_linears(model.model.norm, [model.lm_head]) + for n, p in model.named_parameters(): + assert torch.equal(p.data, snap[n]), f"{spec.name}: re-fusing fused norms changed {n}" + + # (b) identity refold through the full public pipeline + eye_h = torch.eye(spec.hidden, dtype=torch.float64) + eye_d = torch.eye(spec.head_dim, dtype=torch.float64) + rots = fold_rotations(model, R1=eye_h, R2=[eye_d] * spec.n_layers) + for n, p in model.named_parameters(): + assert torch.equal(p.data, snap[n]), f"{spec.name}: identity refold changed {n}" + assert torch.equal(rots["R1"], eye_h), f"{spec.name}: identity R1 not returned" + + +# -------------------------------------------------------------------------------------- +# 4. R.bin round-trip of fold-returned dicts +# -------------------------------------------------------------------------------------- + + +def test_rbin_roundtrip_and_load_gates(): + """The dict fold_rotations returns torch.saves as an R.bin and round-trips bitwise + through RotationSet.load (and through RotationSet.save). load refuses off-manifold + matrices (both gross and drift-sized); orthogonalize=True polar-retracts a drifted + file back onto the manifold (audit < 1e-10, staying near the drifted matrix), and the + retracted set passes fold_rotations' external gate.""" + spec = _CONFIGS[4] # qwen3, 2 layers + rots = fold_rotations(_build(spec), mode="hadamard", seed=2, use_r2=True) + fd, path = tempfile.mkstemp(suffix=".bin") + os.close(fd) + try: + # fold-returned dict -> R.bin -> RotationSet: bitwise + torch.save(dict(rots), path) + rs = RotationSet.load(path) + assert set(rs.rotations) == set(rots) + for k in rots: + assert torch.equal(rs.rotations[k], rots[k]), f"{k} changed in transit" + assert torch.equal(rs.R1, rots["R1"]) + assert set(rs.R2) == _r2_keys(spec.n_layers) + + # RotationSet.save -> load: bitwise again + rs.save(path) + rs2 = RotationSet.load(path) + for k in rots: + assert torch.equal(rs2.rotations[k], rots[k]), f"{k} changed via rs.save" + + # gross off-manifold: refuse + bad = dict(rots) + bad["R1"] = rots["R1"] * 1.01 + torch.save(bad, path) + with pytest.raises(ValueError, match="not orthogonal"): + RotationSet.load(path) + + # drift-sized off-manifold (legacy raw R.bin style): refuse without orthogonalize + torch.manual_seed(3) + drifted = dict(rots) + drifted["R1"] = rots["R1"] + 5e-4 * torch.randn_like(rots["R1"]) + torch.save(drifted, path) + with pytest.raises(ValueError, match="not orthogonal"): + RotationSet.load(path) + + # orthogonalize=True: polar retraction back onto the manifold, near the input + rs3 = RotationSet.load(path, orthogonalize=True) + assert max(rs3.ortho_audit().values()) < 1e-10 + assert (rs3.R1 - drifted["R1"]).abs().max().item() < 1e-2 + # and the retracted set is fold-deployable (passes the external ortho gate) + fold_rotations(_build(spec), R1=rs3.R1, R2=rs3.R2) + finally: + os.unlink(path) + + +# -------------------------------------------------------------------------------------- +# 5. use_r2=False on every config; R2 argument forms are equivalent +# -------------------------------------------------------------------------------------- + + +def test_use_r2_false_all_configs(): + """R1-only fold on every sweep config: returns exactly {'R1'}, keeps fp32 logits + within tolerance, fuses norms to ones, and changes no parameter shape.""" + for spec in _CONFIGS: + tag = f"{spec.name} use_r2=False" + model = _build(spec) + shapes_before = {n: tuple(p.shape) for n, p in model.named_parameters()} + before = _logits(model) + rots = fold_rotations(model, mode=spec.mode, seed=1, use_r2=False) + after = _logits(model) + assert set(rots) == {"R1"}, f"{tag}: keys {set(rots)}" + max_diff = (after - before).abs().max().item() + assert torch.allclose(after, before, rtol=0, atol=ATOL_FP32), ( + f"{tag}: max |delta logit| = {max_diff:.3e} > {ATOL_FP32}" + ) + params_after = dict(model.named_parameters()) + for n, shp in shapes_before.items(): + assert tuple(params_after[n].shape) == shp, f"{tag}: shape of {n} changed" + for layer in model.model.layers: + assert torch.all(layer.input_layernorm.weight == 1), tag + assert torch.all(layer.post_attention_layernorm.weight == 1), tag + assert torch.all(model.model.norm.weight == 1), tag + + +def test_r2_list_dict_int_forms_equivalent(): + """R2 as a layer-ordered list, as a str-keyed dict (R.bin convention) and as an + int-keyed dict produce bitwise-identical folds.""" + for spec in (_CONFIGS[1], _CONFIGS[7]): # llama + qwen3, both multi-form + rots = fold_rotations(_build(spec), mode=spec.mode, seed=4, use_r2=True) + R1 = rots["R1"] + r2_str = {k: v for k, v in rots.items() if k != "R1"} + r2_list = [rots[f"model.layers.{i}.self_attn.R2"] for i in range(spec.n_layers)] + r2_int = {i: rots[f"model.layers.{i}.self_attn.R2"] for i in range(spec.n_layers)} + + m_str, m_list, m_int = _build(spec), _build(spec), _build(spec) + fold_rotations(m_str, R1=R1, R2=r2_str) + fold_rotations(m_list, R1=R1, R2=r2_list) + fold_rotations(m_int, R1=R1, R2=r2_int) + p_str = dict(m_str.named_parameters()) + for other, form in ((m_list, "list"), (m_int, "int-dict")): + for n, p in other.named_parameters(): + assert torch.equal(p.data, p_str[n].data), ( + f"{spec.name}: {n} differs between str-dict and {form} R2 forms" + ) + + +# -------------------------------------------------------------------------------------- +# 6. Orientation oracle: hand-built 1-layer model vs independent manual matmuls +# -------------------------------------------------------------------------------------- + + +def test_orientation_oracle_manual_matmul(): + """Hand-built 1-layer Llama (hidden 64, 4 Q heads, 2 KV heads — GQA 2 — head_dim 32 + != hidden/heads): fold a copy, then verify the orientation algebra directly against + independent fp64 manual matmuls (computed WITHOUT reusing fold's helper functions): + + weight-level — q/k/v readers carry (W * gamma_in) @ R1 (v additionally per-KV-head + rows R2^T W_h), o_proj carries R1^T W with per-Q-head columns @ R2, gate/up carry + (W * gamma_post) @ R1, down_proj carries R1^T W, embed carries E @ R1 (NO gamma), + lm_head carries (W * gamma_final) @ R1; + + seam-level — for random x: readers reproduce the original fused projections from the + rotated stream x@R1; v heads come out R2-rotated; pushing head-mixed values through + the rotated o_proj lands the original output rotated by R1 (same for down_proj and + lm_head); RMSNorm with unit gain commutes with R1 (the identity that makes the whole + fold work); and the two actual models produce equal logits.""" + n_q, n_kv, d, hidden, inter, vocab = 4, 2, 32, 64, 96, 64 + n_rep = n_q // n_kv + + def build(): + torch.manual_seed(4321) + cfg = LlamaConfig( + vocab_size=vocab, + hidden_size=hidden, + intermediate_size=inter, + num_hidden_layers=1, + num_attention_heads=n_q, + num_key_value_heads=n_kv, + head_dim=d, + max_position_embeddings=128, + tie_word_embeddings=False, + attn_implementation="eager", + ) + model = LlamaForCausalLM(cfg).eval() + _randomize_rmsnorm_gains(model) + return model + + model_ref, model_rot = build(), build() + + def f64(t): + return t.data.detach().to(torch.float64).clone() + + L = model_ref.model.layers[0] + A, M = L.self_attn, L.mlp + Wq, Wk, Wv, Wo = ( + f64(A.q_proj.weight), + f64(A.k_proj.weight), + f64(A.v_proj.weight), + f64(A.o_proj.weight), + ) + Wg, Wu, Wd = f64(M.gate_proj.weight), f64(M.up_proj.weight), f64(M.down_proj.weight) + g_in = f64(L.input_layernorm.weight) + g_post = f64(L.post_attention_layernorm.weight) + g_fin = f64(model_ref.model.norm.weight) + E, Wlm = f64(model_ref.model.embed_tokens.weight), f64(model_ref.lm_head.weight) + assert A.q_proj.bias is None and A.o_proj.bias is None and M.down_proj.bias is None + + rots = fold_rotations(model_rot, mode="random", seed=11, use_r2=True) + R1, R2 = rots["R1"], rots["model.layers.0.self_attn.R2"] + Lr = model_rot.model.layers[0] + Ar, Mr = Lr.self_attn, Lr.mlp + Wq_r, Wk_r, Wv_r, Wo_r = ( + f64(Ar.q_proj.weight), + f64(Ar.k_proj.weight), + f64(Ar.v_proj.weight), + f64(Ar.o_proj.weight), + ) + Wg_r, Wu_r, Wd_r = f64(Mr.gate_proj.weight), f64(Mr.up_proj.weight), f64(Mr.down_proj.weight) + E_r, Wlm_r = f64(model_rot.model.embed_tokens.weight), f64(model_rot.lm_head.weight) + + # -------- weight-level oracle (independent per-head loops, pure fp64) -------- + # Stored weights went through fp64 -> fp32 casts, expected values are pure fp64: + # entries are O(0.1), cast noise is O(1e-8), so 1e-6 is tight yet safe; a wrong + # orientation (e.g. R1 @ W instead of W @ R1, or R2 vs R2^T) is an O(0.1) error. + ATOL_W = 1e-6 + + def rot_v_rows(W): # per-KV-head row blocks: W_h <- R2^T @ W_h + out = W.clone() + for h in range(n_kv): + out[h * d : (h + 1) * d, :] = R2.T @ W[h * d : (h + 1) * d, :] + return out + + def rot_o_cols(W): # per-Q-head column blocks: W[:, h*d:(h+1)*d] @ R2 + out = W.clone() + for h in range(n_q): + out[:, h * d : (h + 1) * d] = W[:, h * d : (h + 1) * d] @ R2 + return out + + expected = { + "q_proj": (Wq * g_in) @ R1, # reader (+ input_layernorm gamma) + "k_proj": (Wk * g_in) @ R1, # reader (+ gamma) + "v_proj": rot_v_rows((Wv * g_in) @ R1), # reader (+ gamma) + R2 rows per KV head + "o_proj": rot_o_cols(R1.T @ Wo), # writer + R2 cols per Q head; NO gamma + "gate_proj": (Wg * g_post) @ R1, # reader (+ post_attention gamma) + "up_proj": (Wu * g_post) @ R1, # reader (+ gamma) + "down_proj": R1.T @ Wd, # writer; NO gamma + "embed": E @ R1, # residual writer; NO gamma + "lm_head": (Wlm * g_fin) @ R1, # reader (+ final-norm gamma) + } + got = { + "q_proj": Wq_r, + "k_proj": Wk_r, + "v_proj": Wv_r, + "o_proj": Wo_r, + "gate_proj": Wg_r, + "up_proj": Wu_r, + "down_proj": Wd_r, + "embed": E_r, + "lm_head": Wlm_r, + } + for name in expected: + diff = (got[name] - expected[name]).abs().max().item() + assert diff < ATOL_W, f"weight oracle: {name} max |delta| = {diff:.3e} >= {ATOL_W}" + assert torch.all(Lr.input_layernorm.weight == 1) + assert torch.all(Lr.post_attention_layernorm.weight == 1) + assert torch.all(model_rot.model.norm.weight == 1) + + # -------- seam-level oracle: manual matmuls of the tiny network on random x -------- + ATOL_X = 1e-5 + torch.manual_seed(7) + x = torch.randn(5, hidden, dtype=torch.float64) # residual-stream rows (orig frame) + x_rot = x @ R1 # the SAME stream in the rotated frame + + # readers: rotated weights on the rotated stream reproduce the original projections + for nm, W_fused, W_rot in ( + ("q_proj", Wq * g_in, Wq_r), + ("k_proj", Wk * g_in, Wk_r), + ("gate_proj", Wg * g_post, Wg_r), + ("up_proj", Wu * g_post, Wu_r), + ("lm_head", Wlm * g_fin, Wlm_r), + ): + ref, rot = x @ W_fused.T, x_rot @ W_rot.T + diff = (rot - ref).abs().max().item() + assert diff < ATOL_X, f"seam oracle reader {nm}: max |delta| = {diff:.3e}" + + # v_proj: heads come out rotated by R2 (per KV head) + v_ref = x @ (Wv * g_in).T # [5, n_kv*d] + v_rot = x_rot @ Wv_r.T + for h in range(n_kv): + blk = slice(h * d, (h + 1) * d) + diff = (v_rot[:, blk] - v_ref[:, blk] @ R2).abs().max().item() + assert diff < ATOL_X, f"seam oracle v head {h}: max |delta| = {diff:.3e}" + + # o_proj: attention-mix per Q head (probs are identical in both models because the + # q/k seams above are identities), GQA repeat kv->q via h // n_rep, then the rotated + # o_proj must land the original output rotated by R1. + P = torch.softmax(torch.randn(5, 5, dtype=torch.float64), dim=-1) # stand-in probs + a_ref = torch.cat( + [P @ v_ref[:, (h // n_rep) * d : (h // n_rep + 1) * d] for h in range(n_q)], dim=-1 + ) + a_rot = torch.cat( + [P @ v_rot[:, (h // n_rep) * d : (h // n_rep + 1) * d] for h in range(n_q)], dim=-1 + ) + diff = (a_rot @ Wo_r.T - (a_ref @ Wo.T) @ R1).abs().max().item() + assert diff < ATOL_X, f"seam oracle o_proj: max |delta| = {diff:.3e}" + + # down_proj writer: rotated-frame output = original output @ R1 + m = torch.randn(5, inter, dtype=torch.float64) + diff = (m @ Wd_r.T - (m @ Wd.T) @ R1).abs().max().item() + assert diff < ATOL_X, f"seam oracle down_proj: max |delta| = {diff:.3e}" + + # embed rows land in the rotated frame + diff = (E_r - E @ R1).abs().max().item() + assert diff < ATOL_W, f"seam oracle embed: max |delta| = {diff:.3e}" + + # RMSNorm(unit gain) commutes with the orthogonal R1 — the identity the fold rests on + def rms(v, eps=1e-6): + return v * torch.rsqrt(v.pow(2).mean(-1, keepdim=True) + eps) + + diff = (rms(x_rot) - rms(x) @ R1).abs().max().item() + assert diff < 1e-12, f"RMSNorm rotation-invariance broken: max |delta| = {diff:.3e}" + + # -------- end-to-end: run both actual models on the same ids -------- + ref_logits, rot_logits = _logits(model_ref, vocab), _logits(model_rot, vocab) + max_diff = (rot_logits - ref_logits).abs().max().item() + assert torch.allclose(rot_logits, ref_logits, rtol=0, atol=ATOL_FP32), ( + f"oracle end-to-end: max |delta logit| = {max_diff:.3e} > {ATOL_FP32}" + ) + + +if __name__ == "__main__": + tests = [(n, f) for n, f in sorted(globals().items()) if n.startswith("test_") and callable(f)] + failed = [] + for name, fn in tests: + try: + fn() + print(f"PASS {name}", flush=True) + except Exception: + failed.append(name) + print(f"FAIL {name}", flush=True) + traceback.print_exc() + print( + f"\n{len(tests) - len(failed)}/{len(tests)} tests passed" + + (f"; FAILED: {failed}" if failed else "") + ) + sys.exit(1 if failed else 0) diff --git a/tests/unit/torch/quantization/test_rotation_ext_learner.py b/tests/unit/torch/quantization/test_rotation_ext_learner.py new file mode 100644 index 00000000000..5b9aadbebd1 --- /dev/null +++ b/tests/unit/torch/quantization/test_rotation_ext_learner.py @@ -0,0 +1,656 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Learner-semantics tests for modelopt.torch.quantization.rotation.learn. + +Covers: + 1. steps=0 == fold_rotations seed draws, bitwise, seeds {0, 3, 7}, both archs. + 2. Gradient exclusivity at step 0 for every objective preset (W4A4_G128, + INT8_DEFAULT, int8 per-token-dynamic): grads reach all R params, none reach + model params. + 3. Loss decreases on a repeated batch (25 steps) for each of the 3 objectives. + 4. a_static_scope semantics: "run" scale = monotone running max, "batch" scale + tracks per-batch amax (can decrease); telemetry monotone; both stay orthogonal. + 5. Warm start from a prior RotationSet: step-0 loss continues the donor's final + loss; the ortho gate rejects a corrupted warm start. + 6. Gradient checkpointing: recompute-during-backward still sees the effective + weights (_reparametrize_module spans backward), grads reach R, loss finite. + 7. Objective coverage: activation hooks on exactly 7*n_layers linears; + lm_head/embeddings never hooked and never fake-quantized. + +Plain test_* functions with asserts: collectable by pytest, and also runnable without +it via ``python test_rotation_ext_learner.py`` (the __main__ driver runs every test +function and exits nonzero on any failure). CPU-only, tiny models, seconds per test. +""" + +import os + +os.environ["CUDA_VISIBLE_DEVICES"] = "" # HARD CONSTRAINT: CPU-only, never touch the GPU + +import itertools +import sys +import traceback +import types + +import pytest +import torch +import torch.nn as nn +from torch.nn.utils import stateless +from transformers import LlamaConfig, LlamaForCausalLM, Qwen3Config, Qwen3ForCausalLM + +from modelopt.torch.quantization.rotation import ( + INT8_DEFAULT_OBJECTIVE, + W4A4_G128_OBJECTIVE, + QuantObjective, + fold_rotations, + learn_rotations, +) +from modelopt.torch.quantization.rotation.learn import ( + _ATTN_PROJS, + _MLP_PROJS, + _ActQuantHooks, + _assemble_effective_weights, + _fq_act, +) + +VOCAB = 128 +# Standard tiny fixture (same as test_rotation_{fold,learn}.py): HEAD_DIM deliberately +# decoupled from HIDDEN // num_attention_heads (32 vs 64//4 = 16), like Qwen3-0.6B. +HIDDEN = 64 +HEAD_DIM = 32 +N_LAYERS = 2 + +# Grouped fixture for W4A4_G128_OBJECTIVE: w_group=128 must divide every quantized +# in_features. hidden=128 (q/k/v/gate/up in), heads=4 * head_dim=64 -> o_proj in 256, +# intermediate=2*hidden=256 (down_proj in) — all divisible by 128. head_dim stays +# decoupled (64 != 128//4 = 32). +G_HIDDEN = 128 +G_HEAD_DIM = 64 + +#: The exact three presets of the task: the module's two shipped presets plus the +#: int8 per-token-dynamic recipe. +INT8_PTDYN = QuantObjective( + name="ptdyn", w_bits=8, w_group=None, a_bits=8, a_mode="per_token_dynamic" +) +PRESET_OBJECTIVES = (W4A4_G128_OBJECTIVE, INT8_DEFAULT_OBJECTIVE, INT8_PTDYN) + +# Tiny-model W4A4 for the standard fixture (g=16 divides 64/128/128), as in +# test_rotation_learn.py — used where the objective flavor is not the thing under test. +TINY_W4A4 = QuantObjective( + name="tiny_w4a4", w_bits=4, w_group=16, a_bits=4, a_mode="per_token_dynamic" +) + +# fp32 Cayley iterates drift off the manifold at ~1e-7/step (measured); a handful of +# steps stays well under 1e-5. Post-retraction audits sit at fp64-SVD level (<1e-10). +ORTHO_TOL_FP32_STEPS = 1e-5 + + +def _randomize_rmsnorm_gains(model): + """Non-one RMSNorm gains so norm-fusion effects are not vacuous (fresh HF models + initialize all norm weights to ones).""" + for module in model.modules(): + if type(module).__name__.endswith("RMSNorm"): + module.weight.data = 1.0 + 0.1 * torch.randn_like(module.weight.data) + + +def _build(cfg_cls, model_cls, hidden, head_dim, tie=False): + torch.manual_seed(1234) + cfg = cfg_cls( + vocab_size=VOCAB, + hidden_size=hidden, + intermediate_size=2 * hidden, + num_hidden_layers=N_LAYERS, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=head_dim, + max_position_embeddings=128, + tie_word_embeddings=tie, + attn_implementation="eager", + ) + model = model_cls(cfg).eval() + _randomize_rmsnorm_gains(model) + return model + + +def _tiny_llama(tie=False): + return _build(LlamaConfig, LlamaForCausalLM, HIDDEN, HEAD_DIM, tie) + + +def _tiny_qwen3(tie=False): + return _build(Qwen3Config, Qwen3ForCausalLM, HIDDEN, HEAD_DIM, tie) + + +def _grouped_llama(tie=False): + return _build(LlamaConfig, LlamaForCausalLM, G_HIDDEN, G_HEAD_DIM, tie) + + +def _grouped_qwen3(tie=False): + return _build(Qwen3Config, Qwen3ForCausalLM, G_HIDDEN, G_HEAD_DIM, tie) + + +def _calib_batches(n_batches=2, bs=2, seq=16, seed=7): + torch.manual_seed(seed) + return [torch.randint(0, VOCAB, (bs, seq)) for _ in range(n_batches)] + + +# -------------------------------------------------------------------------------------- +# 1. steps=0 == fold seed draws, bitwise, seeds {0, 3, 7}, both archs +# -------------------------------------------------------------------------------------- + + +def test_steps0_matches_fold_seed_draws_bitwise(): + """learn_rotations(steps=0, seed=s) returns exactly (torch.equal, bitwise) the + matrices fold_rotations(mode, seed=s) draws — for seeds {0, 3, 7} on both archs + (hadamard everywhere; "random" mode additionally checked at seed 0). One RNG + contract: R1 first, then R2 by ascending layer, from the seeded global CPU RNG.""" + for build in (_tiny_llama, _tiny_qwen3): + for seed in (0, 3, 7): + modes = ("hadamard", "random") if seed == 0 else ("hadamard",) + for mode in modes: + init = learn_rotations( + build(), + [], + steps=0, + mode=mode, + objective_cfg=None, + seed=seed, + log_every=0, + ) + folded = fold_rotations(build(), mode=mode, seed=seed, use_r2=True) + assert set(init.rotations) == set(folded), ( + f"{build.__name__} seed={seed} mode={mode}: key sets differ" + ) + for k in folded: + assert torch.equal(init.rotations[k], folded[k]), ( + f"{build.__name__} seed={seed} mode={mode}: {k} not bitwise equal" + ) + + +# -------------------------------------------------------------------------------------- +# 2. Gradient exclusivity at step 0 for every objective preset +# -------------------------------------------------------------------------------------- + + +def _one_training_backward(model, init, objective): + """Mirror one step-0 forward/backward of learn_rotations' loop using the module's + own building blocks (_assemble_effective_weights + _ActQuantHooks + + stateless._reparametrize_module) with the R matrices as fresh leaves, so gradients + are observable from the outside. ``model`` must already be prepared (untied, fused, + frozen) — learn_rotations(steps=0) does exactly that.""" + n_layers = len(model.model.layers) + head_dim = model.config.head_dim + R1 = nn.Parameter(init.R1.to(torch.float32)) + R2s = [ + nn.Parameter(init.rotations[f"model.layers.{i}.self_attn.R2"].to(torch.float32)) + for i in range(n_layers) + ] + sd = dict(model.named_parameters()) + base = { + "model.embed_tokens.weight": sd["model.embed_tokens.weight"].data, + "lm_head.weight": sd["lm_head.weight"].data, + } + for i in range(n_layers): + for proj in _ATTN_PROJS + _MLP_PROJS: + name = f"model.layers.{i}.{proj}.weight" + base[name] = sd[name].data + + hooks = None + if objective is not None and objective.a_bits is not None: + hooks = _ActQuantHooks(objective) + n_hooked = hooks.attach(model) + assert n_hooked == 7 * n_layers, f"hooked {n_hooked}, expected {7 * n_layers}" + torch.manual_seed(42) + ids = torch.randint(0, VOCAB, (2, 16)) + try: + eff = _assemble_effective_weights( + base, + R1, + R2s, + n_layers, + head_dim, + objective, + model.model.embed_tokens.weight.dtype, + ) + with stateless._reparametrize_module(model, eff): + loss = model(input_ids=ids, labels=ids, use_cache=False).loss + loss.backward() + finally: + if hooks is not None: + hooks.remove() + return loss, R1, R2s + + +def test_gradient_exclusivity_every_objective(): + """For each of the 3 objective presets on both archs: one forward/backward at the + seed init gives (i) a finite loss, (ii) a non-None, nonzero, finite gradient on R1 + and on EVERY per-layer R2, and (iii) no gradient on any model parameter.""" + for build in (_grouped_llama, _grouped_qwen3): + for obj in PRESET_OBJECTIVES: + tag = f"{build.__name__}/{obj.name}" + model = build() + init = learn_rotations(model, [], steps=0, objective_cfg=None, seed=0, log_every=0) + loss, R1, R2s = _one_training_backward(model, init, obj) + assert torch.isfinite(loss), f"{tag}: loss not finite: {loss}" + assert R1.grad is not None, f"{tag}: no grad reached R1" + assert torch.isfinite(R1.grad).all(), f"{tag}: R1 grad not finite" + assert R1.grad.abs().max().item() > 0, f"{tag}: R1 grad identically zero" + for i, r2 in enumerate(R2s): + assert r2.grad is not None, f"{tag}: no grad reached R2[{i}]" + assert torch.isfinite(r2.grad).all(), f"{tag}: R2[{i}] grad not finite" + assert r2.grad.abs().max().item() > 0, f"{tag}: R2[{i}] grad zero" + for name, p in model.named_parameters(): + assert p.grad is None, f"{tag}: model param {name} received a gradient" + + +# -------------------------------------------------------------------------------------- +# 3. Loss decreases on a repeated batch for each of the 3 objectives +# -------------------------------------------------------------------------------------- + + +def test_loss_decreases_each_objective(): + """25 Cayley steps overfitting a single repeated batch must reduce the loss for + every preset: min of the last 5 recorded losses < the step-0 loss. (The loop is + deterministic on a repeated batch, so any decrease is real descent, not noise.)""" + for obj in PRESET_OBJECTIVES: + model = _grouped_qwen3() + batch = _calib_batches(n_batches=1, bs=2, seq=32, seed=11) + rs = learn_rotations(model, batch, steps=25, lr=1.0, objective_cfg=obj, seed=0, log_every=0) + losses = [r["loss"] for r in rs.history] + assert len(losses) == 25 + assert all(torch.isfinite(torch.tensor(v)) for v in losses), ( + f"{obj.name}: non-finite loss in {losses}" + ) + assert min(losses[-5:]) < losses[0], ( + f"{obj.name}: no decrease — first {losses[0]:.6f}, last5 {losses[-5:]}" + ) + + +# -------------------------------------------------------------------------------------- +# 4. a_static_scope semantics ("run" vs "batch") +# -------------------------------------------------------------------------------------- + + +class _HookProbe(nn.Module): + """Minimal module whose single Linear is named ``model.layers.0.self_attn.q_proj`` + (the attach filter needs ``.layers.`` in the name + a target suffix), plus one + decoy Linear (``model.head``) that must NOT be hooked.""" + + def __init__(self, in_f=8): + super().__init__() + + class _Attn(nn.Module): + def __init__(self): + super().__init__() + self.q_proj = nn.Linear(in_f, 4, bias=False) + + def forward(self, x): + return self.q_proj(x) + + class _Layer(nn.Module): + def __init__(self): + super().__init__() + self.self_attn = _Attn() + + def forward(self, x): + return self.self_attn(x) + + class _Core(nn.Module): + def __init__(self): + super().__init__() + self.layers = nn.ModuleList([_Layer()]) + # Decoy: an nn.Linear named "model.head" — no ".layers." in its name, + # so the attach filter must skip it. + self.head = nn.Linear(4, 4, bias=False) + + def forward(self, x): + return self.head(self.layers[0](x)) + + self.model = _Core() + + def forward(self, x): + return self.model(x) + + +def test_a_static_scope_run_vs_batch_semantics(): + """Step-resolved semantics of the per-tensor-static activation scale: + + - scope="run": the effective scale is the monotone running max of batch amaxes — + a later small batch is quantized on the stale coarse grid (3/127 -> 0). + - scope="batch": the effective scale tracks each batch's own amax and DECREASES + when a smaller batch arrives — the small batch reconstructs exactly on its grid. + - static_amax telemetry is the monotone non-decreasing observed max in BOTH scopes + (the module records the running max as telemetry either way, by design). + + Verified bitwise against the module's own _fq_act with independently computed + expected scales, via a capture hook registered after the quant hook.""" + key = "model.layers.0.self_attn.q_proj" + qpos = 127.0 + # amax sequence 8 -> 1 -> 5 (decrease then partial recovery); b2/b3 values sit + # exactly on their own amax/127 grids so batch-scope fq is an exact identity. + b1 = torch.tensor([[8.0, -8.0, 4.0, 2.0, 1.0, 0.5, 0.25, 0.0]]) + b2 = torch.tensor([[3.0, -5.0, 1.0, 127.0, -127.0, 2.0, 64.0, 0.0]]) / 127.0 + b3 = torch.tensor([[5.0, -5.0, 2.5, 1.25, 0.5, 0.25, 3.0, 0.0]]) + batches = [b1, b2, b3] + amaxes = [8.0, 1.0, 5.0] + + for scope in ("run", "batch"): + probe = _HookProbe() + cfg = QuantObjective( + name=f"probe_{scope}", + w_bits=None, + w_group=None, + a_bits=8, + a_mode="per_tensor_static", + a_static_scope=scope, + ) + hooks = _ActQuantHooks(cfg) + n = hooks.attach(probe) + assert n == 1, f"scope={scope}: attached {n} hooks, expected 1 (decoy hooked?)" + + captured = [] + lin = probe.model.layers[0].self_attn.q_proj + cap_handle = lin.register_forward_pre_hook( + lambda m, inp: captured.append(inp[0].detach().clone()) + ) + telemetry = [] + with torch.no_grad(): + for x in batches: + probe(x) + telemetry.append(float(hooks.static_amax[key])) + cap_handle.remove() + hooks.remove() + + # Telemetry: monotone non-decreasing running max in BOTH scopes. + assert telemetry == [8.0, 8.0, 8.0], f"scope={scope}: telemetry {telemetry}" + assert all(b >= a for a, b in itertools.pairwise(telemetry)) + + # Effective scale per batch, replicated with the module's own op sequence. + prev = None + for i, (x, cap) in enumerate(zip(batches, captured)): + batch_amax = x.detach().abs().amax() + assert batch_amax.item() == amaxes[i], f"fixture broke: batch {i} amax" + run_max = batch_amax if prev is None else torch.maximum(prev, batch_amax) + prev = run_max + amax = batch_amax if scope == "batch" else run_max + s = (amax / qpos).clamp_min(1e-12) + expected = _fq_act(x, s, 8) + assert torch.equal(cap, expected), ( + f"scope={scope} batch {i}: fq output does not match the " + f"{'per-batch' if scope == 'batch' else 'running-max'} scale " + f"(amax used should be {amax.item():.6f})" + ) + if scope == "run": + # Small batch on the stale coarse grid: 3/127 quantizes to 0 with s=8/127. + assert captured[1][0, 0].item() == 0.0, "run scope: expected 3/127 -> 0" + assert not torch.allclose(captured[1], b2), ( + "run scope: small batch must be distorted by the stale scale" + ) + else: + # Per-batch scale decreased 8/127 -> 1/127: exact-grid identity on b2. + assert torch.allclose(captured[1], b2, rtol=0, atol=1e-9), ( + "batch scope: small batch must reconstruct exactly on its own grid" + ) + + # Both scopes stay orthogonal end-to-end through learn_rotations. + for scope in ("batch", "run"): + obj = QuantObjective( + name=f"int8_{scope}", + w_bits=8, + w_group=None, + a_bits=8, + a_mode="per_tensor_static", + a_static_scope=scope, + ) + rs = learn_rotations( + _tiny_llama(), + _calib_batches(), + steps=4, + lr=1.0, + objective_cfg=obj, + seed=0, + log_every=0, + ) + audit = rs.ortho_audit() + assert max(audit.values()) < 1e-10, f"scope={scope}: post-retraction {audit}" + assert all(r["r1_ortho"] < ORTHO_TOL_FP32_STEPS for r in rs.history), ( + f"scope={scope}: raw drift {[r['r1_ortho'] for r in rs.history]}" + ) + amax = rs.meta.get("static_act_amax", {}) + assert len(amax) == 7 * N_LAYERS and all(v > 0 for v in amax.values()) + + +# -------------------------------------------------------------------------------------- +# 5. Warm start +# -------------------------------------------------------------------------------------- + + +def test_warm_start_continues_donor_loss_and_gate_rejects_corruption(): + """init_rotations from a donor RotationSet: the warm run's step-0 loss (same + repeated batch, identically built model) lands in the ballpark of the donor's + final loss — training progress transfers through save/warm-start — and is closer + to the donor's end than to its start. A deliberately corrupted warm start (R1 + scaled by 1.01, ortho residual ~2e-2 >> 1e-4) is rejected by the ortho gate.""" + batch = _calib_batches(n_batches=1, bs=2, seq=32, seed=11) + donor = learn_rotations( + _tiny_qwen3(), + batch, + steps=12, + lr=1.0, + objective_cfg=TINY_W4A4, + seed=0, + log_every=0, + ) + donor_first = donor.history[0]["loss"] + donor_last = donor.history[-1]["loss"] + assert donor_last < donor_first, ( + f"donor did not descend ({donor_first:.6f} -> {donor_last:.6f}); " + "warm-start continuity check would be vacuous" + ) + + warm = learn_rotations( + _tiny_qwen3(), + batch, + steps=1, + lr=1.0, + objective_cfg=TINY_W4A4, + seed=0, + init_rotations=donor.rotations, + log_every=0, + ) + assert warm.meta["warm_start"] is True + w0 = warm.history[0]["loss"] + # Ballpark: the only differences vs. donor_last are the donor's final near-zero-lr + # Cayley step (cosine lr at step 11/12 is 1.7% of peak), the final polar retraction + # (~1e-6 entry delta), and fp64->fp32 casts. Budget 20% of the donor's total + # improvement, floored at 0.02 absolute. + tol = max(0.02, 0.2 * (donor_first - donor_last)) + assert abs(w0 - donor_last) < tol, ( + f"warm step-0 loss {w0:.6f} not in the donor-final ballpark " + f"{donor_last:.6f} (donor first {donor_first:.6f}, tol {tol:.6f})" + ) + assert abs(w0 - donor_last) < abs(w0 - donor_first), ( + f"warm step-0 loss {w0:.6f} closer to donor start {donor_first:.6f} " + f"than to donor end {donor_last:.6f}" + ) + + # Ortho gate: corrupted warm start must be refused. + bad = dict(donor.rotations) + bad["R1"] = bad["R1"] * 1.01 + with pytest.raises(ValueError, match="not orthogonal"): + learn_rotations( + _tiny_qwen3(), + batch, + steps=1, + objective_cfg=TINY_W4A4, + seed=0, + init_rotations=bad, + log_every=0, + ) + + +# -------------------------------------------------------------------------------------- +# 6. Gradient checkpointing — reparametrize must span backward +# -------------------------------------------------------------------------------------- + + +def test_gradient_checkpointing_grads_still_reach_rotations(): + """gradient_checkpointing_enable() + 3 learn steps: the checkpointed recompute + happens during backward and must still see the effective weights (loss.backward() + runs inside the _reparametrize_module context — the design property). Grads reach + R (the trainer's internal step-0 exclusivity assert would raise otherwise, and the + returned R1 visibly moved from its init), every recorded loss is finite, and the + checkpoint function demonstrably ran (call counter == steps * n_layers). + + transformers' GradientCheckpointingLayer only checkpoints when ``self.training`` + is True, but learn_rotations calls model.eval(); the test force-keeps train mode + via a no-op eval override so the checkpointed path actually executes (the tiny + models have no active dropout, so train-mode forward is deterministic).""" + steps = 3 + model = _tiny_qwen3() + model.gradient_checkpointing_enable() + model.train() + model.eval = types.MethodType(lambda self: self, model) # keep checkpointing active + + calls = {"n": 0} + wrapped_any = False + for m in model.modules(): + f = getattr(m, "_gradient_checkpointing_func", None) + if f is not None and getattr(m, "gradient_checkpointing", False): + + def counting(*a, __f=f, **k): + calls["n"] += 1 + return __f(*a, **k) + + m._gradient_checkpointing_func = counting + wrapped_any = True + assert wrapped_any, "gradient_checkpointing_enable() installed no checkpoint funcs" + + rs = learn_rotations( + model, + _calib_batches(), + steps=steps, + lr=1.0, + objective_cfg=TINY_W4A4, + seed=0, + log_every=0, + ) + assert calls["n"] == steps * N_LAYERS, ( + f"checkpointing did not run as expected: {calls['n']} checkpoint calls, " + f"expected {steps * N_LAYERS} (steps * n_layers)" + ) + assert len(rs.history) == steps + assert all(torch.isfinite(torch.tensor(r["loss"])) for r in rs.history), ( + f"non-finite loss under checkpointing: {[r['loss'] for r in rs.history]}" + ) + # Grads reached R: the returned R1 moved away from the seed-0 init draw. + init = learn_rotations(_tiny_qwen3(), [], steps=0, objective_cfg=None, seed=0, log_every=0) + dmax = (rs.R1 - init.R1).abs().max().item() + assert dmax > 1e-5, f"R1 did not move under checkpointing (max delta {dmax:.3e})" + assert max(rs.ortho_audit().values()) < 1e-10 + + +# -------------------------------------------------------------------------------------- +# 7. Objective coverage — hook targets and never-quantized embeddings/lm_head +# -------------------------------------------------------------------------------------- + + +def test_hook_coverage_and_embeddings_never_quantized(): + """Activation hooks attach to exactly the 7*n_layers decoder-layer linears on both + archs — never lm_head (an nn.Linear!) or embed_tokens — and are fully removed both + by hooks.remove() and after a learn_rotations run. In the effective-weight + assembly, embed_tokens/lm_head are rotated but NEVER weight-fake-quantized, while + decoder linears are.""" + expected = {f"model.layers.{i}.{p}" for i in range(N_LAYERS) for p in _ATTN_PROJS + _MLP_PROJS} + for build in (_tiny_llama, _tiny_qwen3): + model = build() + hooks = _ActQuantHooks(INT8_PTDYN) + n = hooks.attach(model) + assert n == 7 * N_LAYERS, f"{build.__name__}: {n} hooks, expected {7 * N_LAYERS}" + hooked = {name for name, m in model.named_modules() if len(m._forward_pre_hooks) > 0} + assert hooked == expected, ( + f"{build.__name__}: hooked set mismatch: extra {hooked - expected}, " + f"missing {expected - hooked}" + ) + assert "lm_head" not in hooked and "model.embed_tokens" not in hooked + hooks.remove() + assert all(len(m._forward_pre_hooks) == 0 for _, m in model.named_modules()), ( + f"{build.__name__}: hooks leaked after remove()" + ) + + # learn_rotations leaves no hooks behind (removed in its finally block). + model = _tiny_llama() + learn_rotations( + model, + _calib_batches(), + steps=1, + lr=1.0, + objective_cfg=INT8_PTDYN, + seed=0, + log_every=0, + ) + assert all(len(m._forward_pre_hooks) == 0 for _, m in model.named_modules()), ( + "learn_rotations left activation hooks attached" + ) + + # Assembly: embeddings/lm_head rotated but never fake-quantized; decoder linears are. + model = _tiny_qwen3() + init = learn_rotations(model, [], steps=0, objective_cfg=None, seed=0, log_every=0) + n_layers = len(model.model.layers) + R1 = init.R1.to(torch.float32) + R2s = [ + init.rotations[f"model.layers.{i}.self_attn.R2"].to(torch.float32) for i in range(n_layers) + ] + sd = dict(model.named_parameters()) + base = { + "model.embed_tokens.weight": sd["model.embed_tokens.weight"].data, + "lm_head.weight": sd["lm_head.weight"].data, + } + for i in range(n_layers): + for proj in _ATTN_PROJS + _MLP_PROJS: + name = f"model.layers.{i}.{proj}.weight" + base[name] = sd[name].data + # Brutal 2-bit weight quant makes any accidental embed/lm_head quantization obvious. + obj = QuantObjective(name="w2_probe", w_bits=2, w_group=None, a_bits=None) + eff = _assemble_effective_weights( + base, R1, R2s, n_layers, model.config.head_dim, obj, torch.float32 + ) + for name in ("model.embed_tokens.weight", "lm_head.weight"): + plain_rot = (base[name].to(torch.float32) @ R1).to(torch.float32) + assert torch.equal(eff[name], plain_rot), ( + f"{name}: effective weight != plain rotation — it was fake-quantized" + ) + q_name = "model.layers.0.self_attn.q_proj.weight" + plain_q = (base[q_name].to(torch.float32) @ R1).to(torch.float32) + assert not torch.equal(eff[q_name], plain_q), ( + "q_proj effective weight was NOT fake-quantized — objective not applied" + ) + + +if __name__ == "__main__": + tests = [(n, f) for n, f in sorted(globals().items()) if n.startswith("test_") and callable(f)] + failed = [] + for name, fn in tests: + try: + fn() + print(f"PASS {name}", flush=True) + except Exception: + failed.append(name) + print(f"FAIL {name}", flush=True) + traceback.print_exc() + print( + f"\n{len(tests) - len(failed)}/{len(tests)} tests passed" + + (f"; FAILED: {failed}" if failed else "") + ) + sys.exit(1 if failed else 0) diff --git a/tests/unit/torch/quantization/test_rotation_ext_sgdg.py b/tests/unit/torch/quantization/test_rotation_ext_sgdg.py new file mode 100644 index 00000000000..1deb7b2a6ca --- /dev/null +++ b/tests/unit/torch/quantization/test_rotation_ext_sgdg.py @@ -0,0 +1,268 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Extended SGDG numerics tests: manifold invariance, step cap, quirk contracts, retractions. + +Drives the SGDG port (:class:`modelopt.torch.quantization.rotation.SGDG`) through fully +deterministic seeded runs (the stiefel branch draws its stochastic-QR-retraction trigger +from Python's global ``random`` — one ``randint`` per parameter per step, so seeding +``random`` fixes the trajectory) and asserts orthogonality invariance under adversarial +gradients, the ``alpha = min(t, lr)`` step cap, the documented inert-momentum dead-store +quirk, both retraction branches, and the polar-projection nearest-orthogonal property. + +Plain test_* functions with asserts: collectable by pytest, and also runnable without it +via ``python test_rotation_ext_sgdg.py`` (the __main__ driver runs every test function +and exits nonzero on any failure). CPU-only, tiny matrices, seconds per test. +""" + +import os + +os.environ["CUDA_VISIBLE_DEVICES"] = "" # CPU-only unit tests: never claim a GPU + +import random +import sys +import traceback + +import torch + +from modelopt.torch.quantization.rotation import SGDG +from modelopt.torch.quantization.rotation.learn import _polar_project +from modelopt.torch.quantization.rotation.sgdg import _qr_retraction + + +def _ortho_err(R): + """max |R^T R - I| in float64.""" + Rd = R.detach().to(torch.float64) + eye = torch.eye(Rd.shape[0], dtype=torch.float64) + return (Rd.t() @ Rd - eye).abs().max().item() + + +def _seeded_orthogonal_fp32(n, seed): + """Seeded fp32 orthogonal init (fp64 Haar QR cast down — same recipe as the trainer).""" + torch.manual_seed(seed) + q, _ = torch.linalg.qr(torch.randn(n, n, dtype=torch.float64)) + return q.to(torch.float32) + + +def _run_trajectory( + sgdg_cls, + momentum, + steps=20, + n=16, + lr=0.4, + grad_scale=1.0, + param_seed=123, + grad_seed=456, + py_seed=789, +): + """Drive one SGDG (ours or official) through a fully deterministic run. + + Identical seeds -> identical parameter init, identical gradient sequence, identical + global-``random`` state (the stiefel branch consumes exactly one ``randint(1, 101)`` + per parameter per step, so equal ``py_seed`` gives equal retraction triggers). + Returns (list of per-step parameter snapshots, final state momentum_buffer). + """ + P = torch.nn.Parameter(_seeded_orthogonal_fp32(n, param_seed)) + opt = sgdg_cls([P], lr=lr, momentum=momentum, stiefel=True) + gen = torch.Generator().manual_seed(grad_seed) + random.seed(py_seed) + traj = [] + for _ in range(steps): + P.grad = grad_scale * torch.randn(n, n, generator=gen) + opt.step() + traj.append(P.detach().clone()) + buf = opt.state[P]["momentum_buffer"].detach().clone() + return traj, buf + + +# -------------------------------------------------------------------------------------- +# 1. 500-step orthogonality invariance under adversarial gradients +# -------------------------------------------------------------------------------------- + + +def test_500_step_orthogonality_adversarial_grads(): + """500 SGDG steps with adversarial unit-scale randn gradients (far harsher than real + CE gradients): |R^T R - I| < 1e-4 at EVERY step (fp32 Cayley drift + occasional + stochastic QR resets), and the parameter moves far from its init.""" + n, steps = 32, 500 + P = torch.nn.Parameter(_seeded_orthogonal_fp32(n, 0)) + P0 = P.detach().clone() + opt = SGDG([P], lr=1.5, momentum=0.0, stiefel=True) + gen = torch.Generator().manual_seed(1) + random.seed(0) + worst = 0.0 + for step in range(steps): + P.grad = torch.randn(n, n, generator=gen) + opt.step() + err = _ortho_err(P) + worst = max(worst, err) + assert err < 1e-4, f"step {step}: |R^T R - I| = {err:.3e} >= 1e-4" + assert (P.detach() - P0).abs().max().item() > 1e-2, "parameter never moved" + assert worst > 0.0 # fp32 iterates are never exactly on the manifold + + +# -------------------------------------------------------------------------------------- +# 2. Step-cap edge: lr = 100 (alpha = min(t, lr) engages the t cap) +# -------------------------------------------------------------------------------------- + + +def test_step_cap_lr100_no_blowup(): + """lr = 100 with unit-scale gradients: the Cayley step size is capped at + t = 1/(||W||_1 + eps) << lr (verified by replicating the tangent construction for the + first step), so 50 steps neither blow up nor leave the manifold.""" + n = 24 + P = torch.nn.Parameter(_seeded_orthogonal_fp32(n, 2)) + P0 = P.detach().clone() + opt = SGDG([P], lr=100.0, momentum=0.0, stiefel=True) + gen = torch.Generator().manual_seed(3) + random.seed(1) + for step in range(50): + g = torch.randn(n, n, generator=gen) + if step == 0: + # Replicate the step's tangent construction: the cap must engage (t << lr). + unity = P.detach() / P.detach().norm(p=2, dim=1, keepdim=True).add(1e-8) + V = -g.t() + MX = V @ unity + XMX = unity @ MX + XXMX = unity.t() @ XMX + W_hat = MX - 0.5 * XXMX + W = W_hat - W_hat.t() + t = (0.5 * 2 / (W.abs().sum(dim=0).max() + 1e-8)).item() + assert t < 100.0, f"cap never engages: t = {t:.3e} >= lr = 100" + assert t < 1.0, f"expected a tight cap for unit grads, got t = {t:.3e}" + P.grad = g + opt.step() + assert torch.isfinite(P).all(), f"step {step}: non-finite entries at lr=100" + err = _ortho_err(P) + assert err < 1e-4, f"step {step}: |R^T R - I| = {err:.3e} at lr=100" + assert (P.detach() - P0).abs().max().item() > 1e-3, "parameter never moved at lr=100" + + +# -------------------------------------------------------------------------------------- +# 3. Momentum is inert in the stiefel branch (our port, faithfully reproduced quirk) +# -------------------------------------------------------------------------------------- + + +def test_momentum_inert_in_stiefel_branch(): + """momentum=0.0 and momentum=0.9 give BITWISE-identical trajectories in our port: + the official dead store (V rebound to a temp before V.copy_(V_new)) keeps the state + momentum_buffer at exact zeros forever, so the momentum hyperparameter cannot affect + the stiefel update.""" + traj0, buf0 = _run_trajectory(SGDG, 0.0, steps=15, py_seed=11) + traj9, buf9 = _run_trajectory(SGDG, 0.9, steps=15, py_seed=11) + for step, (a, b) in enumerate(zip(traj0, traj9)): + assert torch.equal(a, b), f"step {step}: momentum changed the stiefel trajectory" + assert torch.all(buf0 == 0) and torch.all(buf9 == 0), "momentum_buffer was written" + + +# -------------------------------------------------------------------------------------- +# 4. Forced QR-retraction branch (monkeypatched random.randint) +# -------------------------------------------------------------------------------------- + + +def test_forced_qr_retraction_branch(): + """Force rand_num == 1 (QR retraction) vs rand_num != 1 (plain Cayley) on the same + drifted init and gradient: the forced step re-orthogonalizes the iterate (the control + keeps the drift — the Cayley step preserves, not restores, orthogonality) and the + trajectory measurably changes.""" + n = 24 + torch.manual_seed(5) + q, _ = torch.linalg.qr(torch.randn(n, n, dtype=torch.float64)) + drifted = (q + 5e-3 * torch.randn(n, n, dtype=torch.float64)).to(torch.float32) + g = torch.randn(n, n) + results = {} + orig_randint = random.randint + for label, forced in (("control", 2), ("forced", 1)): + P = torch.nn.Parameter(drifted.clone()) + opt = SGDG([P], lr=0.4, momentum=0.0, stiefel=True) + random.randint = lambda a, b, _v=forced: _v # deterministic branch selection + try: + P.grad = g.clone() + opt.step() + finally: + random.randint = orig_randint + results[label] = P.detach().clone() + assert random.randint is orig_randint # the monkeypatch really was restored + err_forced = _ortho_err(results["forced"]) + err_control = _ortho_err(results["control"]) + assert err_forced < 1e-4, f"forced retraction left the iterate off-manifold: {err_forced:.3e}" + assert err_control > 1e-4, ( + f"contrast broken: control step already orthogonal ({err_control:.3e}) — " + "the drifted init did not drift" + ) + diff = (results["forced"] - results["control"]).abs().max().item() + assert not torch.equal(results["forced"], results["control"]) + assert diff > 1e-5, f"forcing the retraction barely changed the step ({diff:.3e})" + + +# -------------------------------------------------------------------------------------- +# 5. Polar retraction: nearest-orthogonal sanity +# -------------------------------------------------------------------------------------- + + +def test_polar_project_nearest_orthogonal(): + """For 20 random near-orthogonal drifted matrices, _polar_project returns a matrix + that is (a) orthogonal to 1e-12 in BOTH residual forms, and (b) strictly closer to + the input in Frobenius norm than any of 50 random orthogonal probes AND than the QR + retraction of the same input (the polar factor is THE nearest orthogonal matrix).""" + n = 24 + gen = torch.Generator().manual_seed(6) + eye = torch.eye(n, dtype=torch.float64) + for trial in range(20): + q, _ = torch.linalg.qr(torch.randn(n, n, generator=gen, dtype=torch.float64)) + A = q + 1e-2 * torch.randn(n, n, generator=gen, dtype=torch.float64) + P = _polar_project(A) + # (a) orthogonality at fp64-SVD level, both basis-dependent residual forms + rtr = (P.t() @ P - eye).abs().max().item() + rrt = (P @ P.t() - eye).abs().max().item() + assert rtr < 1e-12, f"trial {trial}: |P^T P - I| = {rtr:.3e}" + assert rrt < 1e-12, f"trial {trial}: |P P^T - I| = {rrt:.3e}" + d_polar = (A - P).norm().item() + assert d_polar > 0.0 # the input really was off the manifold + # (b) vs the QR retraction of the same input (clone: _qr_retraction mutates + # its argument in place via t_()) + Q_qr = _qr_retraction(A.clone()) + qr_ortho = (Q_qr.t() @ Q_qr - eye).abs().max().item() + assert qr_ortho < 1e-10, f"trial {trial}: QR retraction not orthogonal" + d_qr = (A - Q_qr).norm().item() + assert d_polar < d_qr, ( + f"trial {trial}: polar ({d_polar:.6e}) not closer than QR ({d_qr:.6e})" + ) + # (b) vs 50 random orthogonal probes + for probe_i in range(50): + probe, _ = torch.linalg.qr(torch.randn(n, n, generator=gen, dtype=torch.float64)) + d_probe = (A - probe).norm().item() + assert d_polar < d_probe, ( + f"trial {trial} probe {probe_i}: polar ({d_polar:.6e}) not closer " + f"than a random probe ({d_probe:.6e})" + ) + + +if __name__ == "__main__": + tests = [(n, f) for n, f in sorted(globals().items()) if n.startswith("test_") and callable(f)] + failed = [] + for name, fn in tests: + try: + fn() + print(f"PASS {name}", flush=True) + except Exception: + failed.append(name) + print(f"FAIL {name}", flush=True) + traceback.print_exc() + print( + f"\n{len(tests) - len(failed)}/{len(tests)} tests passed" + + (f"; FAILED: {failed}" if failed else "") + ) + sys.exit(1 if failed else 0) diff --git a/tests/unit/torch/quantization/test_rotation_fold.py b/tests/unit/torch/quantization/test_rotation_fold.py new file mode 100644 index 00000000000..b0f24964c93 --- /dev/null +++ b/tests/unit/torch/quantization/test_rotation_fold.py @@ -0,0 +1,259 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for modelopt.torch.quantization.rotation.fold_rotations (offline R1+R2 folding). + +Plain test_* functions with asserts: collectable by pytest, and also runnable without it +via ``python test_rotation_fold.py`` (the __main__ driver runs every test function and +exits nonzero on any failure). +""" + +import sys +import traceback + +import pytest +import torch +from transformers import LlamaConfig, LlamaForCausalLM, Qwen3Config, Qwen3ForCausalLM + +from modelopt.torch.quantization.rotation import fold_rotations + +VOCAB = 128 +HIDDEN = 64 +# Deliberately decoupled: HEAD_DIM != HIDDEN // num_attention_heads (32 vs 64//4 = 16), like +# Qwen3-0.6B (128 vs 1024//16 = 64). A coincident config would let a regression of the +# head_dim resolution to the wrong fallback formula pass every test (the R2 shape check in +# test_returned_rotations_orthonormal is the only one that can catch it, and only if the two +# values differ). This also makes o_proj non-square ([64, 128]), covering that trap too. +HEAD_DIM = 32 +N_LAYERS = 2 + +# FP-equivalence tolerance for fp32 round-trip: fold_rotations does all math in float64 and +# error enters only through the float64 -> float32 cast of each rewritten weight (<= 2^-24 +# relative per weight), propagated through 2 decoder layers of an exactly-equivalent +# reparametrization. Measured max |delta logit| on these models is ~3e-7 against logits of +# magnitude ~0.6; atol=1e-4 keeps >100x headroom (also for deeper/bigger real models where +# the cast noise accumulates) while still catching any real transform bug — a wrong rotation +# orientation or a missed norm fusion perturbs logits by O(1). +ATOL_FP32 = 1e-4 + + +def _randomize_rmsnorm_gains(model): + """Set every RMSNorm gain to a random non-one value. Fresh HF models initialize all norm + weights to ones, which would make both the norm-fusion math and the fused-to-ones / + q_norm-untouched checks vacuous.""" + for module in model.modules(): + if type(module).__name__.endswith("RMSNorm"): + module.weight.data = 1.0 + 0.1 * torch.randn_like(module.weight.data) + + +def _tiny_llama(tie=False): + torch.manual_seed(1234) + cfg = LlamaConfig( + vocab_size=VOCAB, + hidden_size=HIDDEN, + intermediate_size=2 * HIDDEN, + num_hidden_layers=N_LAYERS, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=HEAD_DIM, + max_position_embeddings=128, + tie_word_embeddings=tie, + attn_implementation="eager", + ) + model = LlamaForCausalLM(cfg).eval() + _randomize_rmsnorm_gains(model) + return model + + +def _tiny_qwen3(tie=False): + torch.manual_seed(1234) + cfg = Qwen3Config( + vocab_size=VOCAB, + hidden_size=HIDDEN, + intermediate_size=2 * HIDDEN, + num_hidden_layers=N_LAYERS, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=HEAD_DIM, + max_position_embeddings=128, + tie_word_embeddings=tie, + attn_implementation="eager", + ) + model = Qwen3ForCausalLM(cfg).eval() + _randomize_rmsnorm_gains(model) + return model + + +def _logits(model): + torch.manual_seed(99) + ids = torch.randint(0, VOCAB, (2, 8)) + with torch.no_grad(): + return model(ids).logits + + +def test_fp_equivalence(): + """(a) Logits before vs after fold agree within the fp32 round-trip tolerance.""" + for build in (_tiny_llama, _tiny_qwen3): + model = build() + before = _logits(model) + fold_rotations(model, mode="hadamard", seed=0, use_r2=True) + after = _logits(model) + max_diff = (after - before).abs().max().item() + assert torch.allclose(after, before, rtol=0, atol=ATOL_FP32), ( + f"{build.__name__}: max |delta logit| = {max_diff:.3e} > {ATOL_FP32}" + ) + + +def test_fused_norm_weights_are_ones(): + """(b) All fused RMSNorm gains (started random) are exactly ones after fold.""" + for build in (_tiny_llama, _tiny_qwen3): + model = build() + fold_rotations(model) + for layer in model.model.layers: + assert torch.all(layer.input_layernorm.weight == 1) + assert torch.all(layer.post_attention_layernorm.weight == 1) + assert torch.all(model.model.norm.weight == 1) + + +def test_qwen3_qk_norm_bitwise_untouched(): + """(c) Qwen3 per-head q_norm/k_norm are bitwise identical after fold.""" + model = _tiny_qwen3() + before = { + n: p.data.clone() for n, p in model.named_parameters() if "q_norm" in n or "k_norm" in n + } + assert len(before) == 2 * N_LAYERS + # Sanity: gains were randomized, so "untouched" is not trivially "still ones". + assert all(not torch.all(p == 1) for p in before.values()) + fold_rotations(model) + for n, p in model.named_parameters(): + if n in before: + assert torch.equal(p.data, before[n]), f"{n} changed" + + +def test_tied_embeddings_untied_and_lm_head_correct(): + """(d) Tied embeddings: config flag off after fold, storage untied, and lm_head equals + the final-norm-fused + R1-rotated copy of the original shared weight (while embed_tokens + is rotated WITHOUT the norm gain).""" + model = _tiny_qwen3(tie=True) + embed = model.model.embed_tokens + assert model.lm_head.weight.data_ptr() == embed.weight.data_ptr() # really tied + + before = _logits(model) + shared = embed.weight.data.clone() + gamma_final = model.model.norm.weight.data.clone() + rotations = fold_rotations(model) + + assert model.config.tie_word_embeddings is False + assert model.lm_head.weight.data_ptr() != embed.weight.data_ptr() + + # Replicate fold's exact op sequence (fp64 fuse -> fp32 cast -> fp64 rotate -> fp32). + fused = (shared.double() * gamma_final.double()).to(shared.dtype) + expected_head = (fused.double() @ rotations["R1"]).to(shared.dtype) + expected_embed = (shared.double() @ rotations["R1"]).to(shared.dtype) + assert torch.allclose(model.lm_head.weight.data, expected_head, rtol=0, atol=1e-7) + assert torch.allclose(embed.weight.data, expected_embed, rtol=0, atol=1e-7) + + after = _logits(model) + assert torch.allclose(after, before, rtol=0, atol=ATOL_FP32) + + +def test_returned_rotations_orthonormal(): + """(e) Returned R1 + per-layer R2 are float64 CPU and orthonormal to < 1e-10.""" + model = _tiny_qwen3() + rotations = fold_rotations(model) + expected_keys = {"R1"} | {f"model.layers.{i}.self_attn.R2" for i in range(N_LAYERS)} + assert set(rotations) == expected_keys + for name, mat in rotations.items(): + size = HIDDEN if name == "R1" else HEAD_DIM + assert mat.dtype == torch.float64 and mat.device.type == "cpu" + assert mat.shape == (size, size) + err = (mat @ mat.T - torch.eye(size, dtype=torch.float64)).abs().max().item() + assert err < 1e-10, f"{name}: max |R R^T - I| = {err:.3e}" + + +def test_r2_actually_applied_to_v_and_o_weights(): + """(g) R2 is really folded into the weights, not just returned. R2 folding is a + functional identity, so fp-equivalence (test a) can never catch a silent no-op in + _rotate_v_proj_r2/_rotate_o_proj_r2. Fold two identically-built models with the same + seed, one with use_r2=False (R1 is drawn before any R2, so R1 is identical) and one + with use_r2=True, then check the use_r2=True weights equal the use_r2=False weights + with the returned R2 applied: v_proj rows R2^T per KV-head block, o_proj columns @ R2 + per Q-head block (replicating fold's exact fp64 op sequence).""" + for build in (_tiny_llama, _tiny_qwen3): + model_no_r2 = build() + model_r2 = build() # same construction seed -> bitwise-identical weights + rot_no_r2 = fold_rotations(model_no_r2, mode="hadamard", seed=0, use_r2=False) + rotations = fold_rotations(model_r2, mode="hadamard", seed=0, use_r2=True) + assert set(rot_no_r2) == {"R1"} + assert torch.equal(rot_no_r2["R1"], rotations["R1"]) + + for idx in range(N_LAYERS): + R2 = rotations[f"model.layers.{idx}.self_attn.R2"] + attn0 = model_no_r2.model.layers[idx].self_attn + attn1 = model_r2.model.layers[idx].self_attn + + # v_proj: per-KV-head row blocks W_h <- R2^T @ W_h (as W^T blocks @ R2). + v0 = attn0.v_proj.weight.data + out_f, in_f = v0.shape + vt = v0.to(torch.float64).t() + vt = (vt.reshape(in_f, out_f // HEAD_DIM, HEAD_DIM) @ R2).reshape(in_f, out_f) + expected_v = vt.t().contiguous().to(v0.dtype) + got_v = attn1.v_proj.weight.data + assert not torch.equal(got_v, v0), f"layer {idx}: v_proj unchanged by use_r2=True" + assert torch.allclose(got_v, expected_v, rtol=0, atol=1e-7), ( + f"layer {idx}: v_proj does not carry the returned R2" + ) + + # o_proj: per-Q-head column blocks W[:, h*d:(h+1)*d] <- W[:, h*d:(h+1)*d] @ R2. + o0 = attn0.o_proj.weight.data + out_f, in_f = o0.shape + expected_o = ( + (o0.to(torch.float64).reshape(out_f, in_f // HEAD_DIM, HEAD_DIM) @ R2) + .reshape(out_f, in_f) + .to(o0.dtype) + ) + got_o = attn1.o_proj.weight.data + assert not torch.equal(got_o, o0), f"layer {idx}: o_proj unchanged by use_r2=True" + assert torch.allclose(got_o, expected_o, rtol=0, atol=1e-7), ( + f"layer {idx}: o_proj does not carry the returned R2" + ) + + +def test_unsupported_arch_raises(): + """(f) A model class outside the registry raises NotImplementedError.""" + + class NotADecoderLM(torch.nn.Module): + pass + + with pytest.raises(NotImplementedError, match="NotADecoderLM"): + fold_rotations(NotADecoderLM()) + + +if __name__ == "__main__": + tests = [(n, f) for n, f in sorted(globals().items()) if n.startswith("test_") and callable(f)] + failed = [] + for name, fn in tests: + try: + fn() + print(f"PASS {name}", flush=True) + except Exception: + failed.append(name) + print(f"FAIL {name}", flush=True) + traceback.print_exc() + print( + f"\n{len(tests) - len(failed)}/{len(tests)} tests passed" + + (f"; FAILED: {failed}" if failed else "") + ) + sys.exit(1 if failed else 0) diff --git a/tests/unit/torch/quantization/test_rotation_kd.py b/tests/unit/torch/quantization/test_rotation_kd.py new file mode 100644 index 00000000000..96cf758d646 --- /dev/null +++ b/tests/unit/torch/quantization/test_rotation_kd.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the KD rotation objective: learn_rotations(teacher=...).""" + +import copy +import sys +import traceback + +import torch +from transformers import LlamaConfig, LlamaForCausalLM + +from modelopt.torch.quantization.rotation import QuantObjective, learn_rotations + +VOCAB = 128 +HIDDEN = 64 +HEAD_DIM = 32 + +TINY_W4A4 = QuantObjective( + name="tiny_w4a4", w_bits=4, w_group=16, a_bits=4, a_mode="per_token_dynamic" +) + + +def _tiny_llama(): + torch.manual_seed(1234) + cfg = LlamaConfig( + vocab_size=VOCAB, + hidden_size=HIDDEN, + intermediate_size=2 * HIDDEN, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=HEAD_DIM, + max_position_embeddings=128, + tie_word_embeddings=False, + attn_implementation="eager", + ) + model = LlamaForCausalLM(cfg).eval() + for module in model.modules(): + if type(module).__name__.endswith("RMSNorm"): + module.weight.data = 1.0 + 0.1 * torch.randn_like(module.weight.data) + return model + + +def _batches(n=1, bs=2, seq=16, seed=7): + torch.manual_seed(seed) + return [torch.randint(0, VOCAB, (bs, seq)) for _ in range(n)] + + +def test_kd_wiring_no_quant_kl_vanishes(): + """With ALL quantizers off and teacher = an exact functional copy, the rotated + student computes the teacher's function -> KL term ~ 0 and the KD loss equals + (1 - kd_alpha) * CE of the plain run at step 0 (lr=0 isolates step-0).""" + batches = _batches() + ref = learn_rotations( + _tiny_llama(), batches, steps=1, lr=0.0, objective_cfg=None, seed=3, log_every=0 + ) + model = _tiny_llama() + teacher = copy.deepcopy(model) + kd = learn_rotations( + model, + batches, + steps=1, + lr=0.0, + objective_cfg=None, + seed=3, + log_every=0, + teacher=teacher, + kd_alpha=0.5, + kd_temp=2.0, + ) + ce_ref = ref.history[0]["loss"] + got = kd.history[0]["loss"] + want = 0.5 * ce_ref # (1-a)*CE + a*T^2*KL, KL ~ 0 + assert abs(got - want) < 5e-3 * max(1.0, abs(want)), ( + f"KD step-0 loss {got} != (1-a)*CE = {want} (CE {ce_ref}) — wiring or a " + "non-vanishing KL where the student equals the teacher" + ) + assert kd.meta["kd"] == {"alpha": 0.5, "temp": 2.0} + assert ref.meta["kd"] is None + + +def test_kd_objective_trains_and_stays_orthogonal(): + """The KD path is live and Cayley-safe: rotations MOVE, stay orthogonal, and the KD + term is genuinely nonzero against a perturbed teacher. + + Deliberately does NOT assert "final loss < first loss": at this toy scale the + trajectory is noise (plain CE on the same recipe also ends higher than it starts), so + such an assertion passes or fails on seed luck rather than on training behavior. + """ + torch.manual_seed(11) + batch = _batches(n=1) + model = _tiny_llama() + teacher = copy.deepcopy(model) + for p in teacher.parameters(): # a teacher that genuinely differs -> KL > 0 + p.data.add_(0.05 * torch.randn_like(p.data)) + teacher.eval().requires_grad_(False) + + ref = learn_rotations( + _tiny_llama(), batch, steps=0, lr=0.5, objective_cfg=TINY_W4A4, seed=5, log_every=0 + ) + rs = learn_rotations( + model, + batch, + steps=10, + lr=0.5, + objective_cfg=TINY_W4A4, + seed=5, + log_every=0, + teacher=teacher, + ) + # (a) training actually moved the rotations away from the shared seeded init + moved = (rs.R1 - ref.R1).abs().max().item() + assert moved > 1e-4, f"rotations did not move under the KD objective (max delta {moved:.2e})" + # (b) every iterate is still on the manifold after the Cayley steps + retraction + assert max(rs.ortho_audit().values()) < 1e-4 + # (c) the KD term is live: an alpha=1 run against this teacher has a positive loss + kd_only = learn_rotations( + _tiny_llama(), + batch, + steps=1, + lr=0.0, + objective_cfg=TINY_W4A4, + seed=5, + log_every=0, + teacher=teacher, + kd_alpha=1.0, + ) + assert kd_only.history[0]["loss"] > 0.0, "KD term is identically zero (KD path is dead)" + + +def test_teacher_untouched(): + model = _tiny_llama() + teacher = copy.deepcopy(model) + before = {n: p.detach().clone() for n, p in teacher.named_parameters()} + learn_rotations( + model, + _batches(), + steps=3, + lr=0.5, + objective_cfg=TINY_W4A4, + seed=5, + log_every=0, + teacher=teacher, + ) + for n, p in teacher.named_parameters(): + assert torch.equal(p.detach(), before[n]), f"teacher param {n} changed" + + +if __name__ == "__main__": + failures = 0 + for name, fn in sorted(globals().items()): + if name.startswith("test_") and callable(fn): + try: + fn() + print(f"PASS {name}") + except Exception: + failures += 1 + print(f"FAIL {name}") + traceback.print_exc() + sys.exit(1 if failures else 0) diff --git a/tests/unit/torch/quantization/test_rotation_learn.py b/tests/unit/torch/quantization/test_rotation_learn.py new file mode 100644 index 00000000000..fca01dcbc01 --- /dev/null +++ b/tests/unit/torch/quantization/test_rotation_learn.py @@ -0,0 +1,481 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for modelopt.torch.quantization.rotation.learn (Cayley-SGD learned R1/R2) and the +external-matrix path of fold_rotations. + +Plain test_* functions with asserts: collectable by pytest, and also runnable without it +via ``python test_rotation_learn.py`` (the __main__ driver runs every test function and +exits nonzero on any failure). CPU-only, tiny models, seconds per test. +""" + +import os +import sys +import tempfile +import traceback + +import pytest +import torch +from transformers import LlamaConfig, LlamaForCausalLM, Qwen3Config, Qwen3ForCausalLM + +from modelopt.torch.quantization.rotation import ( + INT8_DEFAULT_OBJECTIVE, + SGDG, + QuantObjective, + RotationSet, + fold_rotations, + learn_rotations, +) + +VOCAB = 128 +HIDDEN = 64 +# Same deliberate decoupling as test_rotation_fold.py: HEAD_DIM != HIDDEN // +# num_attention_heads (32 vs 16), like Qwen3-0.6B — a coincident config would let a +# head_dim-resolution regression pass silently. +HEAD_DIM = 32 +N_LAYERS = 2 + +# Tiny-model objective: per-group weights must divide every in_features +# ({64 (hidden), 128 (num_q_heads*head_dim, o_proj), 128 (2*hidden, down_proj)}), so g=16. +TINY_W4A4 = QuantObjective( + name="tiny_w4a4", w_bits=4, w_group=16, a_bits=4, a_mode="per_token_dynamic" +) + +# Orthonormality gate for fp32 Cayley iterates after a handful of steps: each step +# perturbs |R^T R - I| at the fp32-rounding scale (~1e-7 per step, measured); 1e-5 keeps +# margin while still catching any real manifold-departure bug (plain SGD drifts to O(lr)). +ORTHO_TOL_FP32_STEPS = 1e-5 + + +def _randomize_rmsnorm_gains(model): + for module in model.modules(): + if type(module).__name__.endswith("RMSNorm"): + module.weight.data = 1.0 + 0.1 * torch.randn_like(module.weight.data) + + +def _tiny_llama(tie=False): + torch.manual_seed(1234) + cfg = LlamaConfig( + vocab_size=VOCAB, + hidden_size=HIDDEN, + intermediate_size=2 * HIDDEN, + num_hidden_layers=N_LAYERS, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=HEAD_DIM, + max_position_embeddings=128, + tie_word_embeddings=tie, + attn_implementation="eager", + ) + model = LlamaForCausalLM(cfg).eval() + _randomize_rmsnorm_gains(model) + return model + + +def _tiny_qwen3(tie=False): + torch.manual_seed(1234) + cfg = Qwen3Config( + vocab_size=VOCAB, + hidden_size=HIDDEN, + intermediate_size=2 * HIDDEN, + num_hidden_layers=N_LAYERS, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=HEAD_DIM, + max_position_embeddings=128, + tie_word_embeddings=tie, + attn_implementation="eager", + ) + model = Qwen3ForCausalLM(cfg).eval() + _randomize_rmsnorm_gains(model) + return model + + +def _calib_batches(n_batches=2, bs=2, seq=16, seed=7): + torch.manual_seed(seed) + return [torch.randint(0, VOCAB, (bs, seq)) for _ in range(n_batches)] + + +def _logits(model): + torch.manual_seed(99) + ids = torch.randint(0, VOCAB, (2, 8)) + with torch.no_grad(): + return model(ids).logits + + +def _ortho_err(R): + Rd = R.detach().to(torch.float64) + eye = torch.eye(Rd.shape[0], dtype=torch.float64) + return (Rd.t() @ Rd - eye).abs().max().item() + + +# -------------------------------------------------------------------------------------- +# 1. The Cayley step itself preserves orthonormality (optimizer in isolation) +# -------------------------------------------------------------------------------------- + + +def test_sgdg_cayley_preserves_orthonormality(): + """10 SGDG steps on a square Stiefel parameter with adversarially LARGE random + gradients (unit-scale randn — far harsher than real CE gradients): every iterate stays + orthonormal at trained-rotation tolerance, the parameter actually moves, and — the + contrast that makes the gate meaningful — a single plain-SGD step with the same + lr/gradient leaves the manifold by >3 orders of magnitude more. Per-step Cayley drift + is fp32-rounding + 5-iteration fixed-point truncation (~1e-6/step here; a full-scale + reference run measures 4.6e-5 after 150 steps on a 2048-dim R1).""" + torch.manual_seed(0) + q, _ = torch.linalg.qr(torch.randn(32, 32, dtype=torch.float64)) + P = torch.nn.Parameter(q.to(torch.float32)) + P0 = P.detach().clone() + opt = SGDG([P], lr=0.5, stiefel=True) + gen = torch.Generator().manual_seed(1) + G0 = None + for step in range(10): + G = torch.randn(32, 32, generator=gen) + G0 = G if G0 is None else G0 + loss = (P * G).sum() + opt.zero_grad(set_to_none=True) + loss.backward() + opt.step() + err = _ortho_err(P) + assert err < 5e-5, f"step {step}: |P^T P - I| = {err:.3e}" + assert (P.detach() - P0).abs().max().item() > 1e-3, "parameter never moved" + # Contrast: one EUCLIDEAN SGD step (same lr, same first gradient) departs the + # manifold by O(lr * ||G||) — the Cayley step is what preserves it. + err_sgd = _ortho_err(P0 - 0.5 * G0) + assert err_sgd > 1e-1, f"contrast broken: plain-SGD ortho err only {err_sgd:.3e}" + + +# -------------------------------------------------------------------------------------- +# 2. learn_rotations end-to-end on a tiny model +# -------------------------------------------------------------------------------------- + + +def test_learn_rotations_orthonormal_and_moved(): + """A short quantized-objective run returns the full fold-convention key set, float64 + CPU matrices, all orthonormal after training, visibly moved from the init draws, and a + complete per-step history.""" + model = _tiny_qwen3() + steps = 6 + rs = learn_rotations( + model, + _calib_batches(), + steps=steps, + lr=1.0, + objective_cfg=TINY_W4A4, + seed=0, + log_every=0, + ) + expected = {"R1"} | {f"model.layers.{i}.self_attn.R2" for i in range(N_LAYERS)} + assert set(rs.rotations) == expected + for name, R in rs.rotations.items(): + size = HIDDEN if name == "R1" else HEAD_DIM + assert R.dtype == torch.float64 and R.device.type == "cpu" + assert R.shape == (size, size) + # Trained sets are polar-retracted on return: BOTH residual forms at fp64 SVD level. + audit = rs.ortho_audit() + assert max(audit.values()) < 1e-10, f"ortho audit (post-retraction): {audit}" + assert len(rs.history) == steps + assert all(torch.isfinite(torch.tensor(r["loss"])) for r in rs.history) + # Moved from init: compare against the seed-0 draws (== steps=0 output). + init = learn_rotations( + _tiny_qwen3(), + _calib_batches(), + steps=0, + objective_cfg=TINY_W4A4, + seed=0, + log_every=0, + ) + dmax = (rs.R1 - init.R1).abs().max().item() + assert dmax > 1e-4, f"R1 did not move from its init (max delta {dmax:.3e})" + + +def test_init_matches_fold_seed_draws(): + """steps=0 returns exactly (bitwise) the matrices fold_rotations draws for the same + seed/mode — the trainer's init and the validated fold path share one RNG contract.""" + init = learn_rotations( + _tiny_qwen3(), _calib_batches(), steps=0, objective_cfg=None, seed=5, log_every=0 + ) + folded = fold_rotations(_tiny_qwen3(), mode="hadamard", seed=5, use_r2=True) + assert set(init.rotations) == set(folded) + for k in folded: + assert torch.equal(init.rotations[k], folded[k]), f"{k}: init != fold draw" + + +def test_loss_decreases_tiny_overfit(): + """Strict loss decrease on a single repeated batch (overfit): with the fake-quant + objective on, CE depends on R through the quantization error, and a few Cayley steps + must reduce it below the step-0 value.""" + model = _tiny_qwen3() + batch = _calib_batches(n_batches=1, bs=2, seq=32, seed=11) + rs = learn_rotations( + model, batch, steps=25, lr=1.0, objective_cfg=TINY_W4A4, seed=0, log_every=0 + ) + losses = [r["loss"] for r in rs.history] + assert min(losses[-5:]) < losses[0], ( + f"no strict decrease: first {losses[0]:.6f}, last5 {losses[-5:]}" + ) + assert rs.history[-1]["r1_ortho"] < ORTHO_TOL_FP32_STEPS + + +def test_learned_fold_fp_equivalence(): + """Learned R fed through fold_rotations(R1=..., R2=...) keeps fp-equivalence with + quantization off. The tolerance budgets the trained matrices' manifold drift (~1e-6 + after 6 fp32 steps) on O(1) logits; a wrong orientation or a missed fusion is O(1).""" + rs = learn_rotations( + _tiny_qwen3(), + _calib_batches(), + steps=6, + lr=1.0, + objective_cfg=TINY_W4A4, + seed=0, + log_every=0, + ) + model = _tiny_qwen3() # fresh, un-mutated model + before = _logits(model) + applied = fold_rotations(model, R1=rs.R1, R2=rs.R2) + after = _logits(model) + max_diff = (after - before).abs().max().item() + assert torch.allclose(after, before, rtol=0, atol=1e-3), ( + f"max |delta logit| = {max_diff:.3e} > 1e-3" + ) + # The returned dict is the applied (float64) matrices — bitwise the learned ones. + for k, v in rs.rotations.items(): + assert torch.equal(applied[k], v), f"{k}: applied != learned" + + +def test_final_retraction_closes_rrt_gap(): + """The returned matrices are the polar retraction of the raw fp32 iterates: meta + records per-matrix raw drift in BOTH residual forms plus the entry-wise projection + distance; the post-retraction audit is at fp64-SVD level, far below the raw drift. + (Motivating field measurement: raw 150-step R1 passes R^T R at ~5e-5 but sits at + ~1e-3 in the R R^T form the fold consumes — basis-dependent max-entry residuals.)""" + rs = learn_rotations( + _tiny_qwen3(), + _calib_batches(), + steps=8, + lr=1.5, + objective_cfg=TINY_W4A4, + seed=0, + log_every=0, + ) + log = rs.meta["final_retraction"] + assert set(log) == set(rs.rotations) + audit = rs.ortho_audit() + for k, rec in log.items(): + assert rec["raw_rtr"] > 0 and rec["raw_rrt"] > 0 and rec["delta_max"] > 0 + # projection moved entries on the order of the raw drift, not more than ~its size + assert rec["delta_max"] < 10 * max(rec["raw_rtr"], rec["raw_rrt"]) + # retraction actually closed the residual: orders of magnitude below raw drift + assert audit[k] < 1e-10 < rec["raw_rrt"] + + +def test_qwen3_qk_norm_bitwise_untouched_by_learn(): + """Qwen3 per-head q_norm/k_norm (head-space, post-projection) are bitwise identical + after learn_rotations — the arch spec excludes them from fusion and rotation.""" + model = _tiny_qwen3() + before = { + n: p.data.clone() for n, p in model.named_parameters() if "q_norm" in n or "k_norm" in n + } + assert len(before) == 2 * N_LAYERS + assert all(not torch.all(p == 1) for p in before.values()) # gains were randomized + learn_rotations( + model, + _calib_batches(), + steps=3, + lr=1.0, + objective_cfg=TINY_W4A4, + seed=0, + log_every=0, + ) + for n, p in model.named_parameters(): + if n in before: + assert torch.equal(p.data, before[n]), f"{n} changed" + + +def test_int8_static_objective_smoke(): + """INT8_DEFAULT_CFG axes (per-out-channel W8 + per-tensor static A8): trains on the + Llama tiny model in BOTH static-scope variants (batch = scale tracks current R, + run = monotone running max), stays orthonormal, and records a positive static amax + per target linear.""" + for scope in ("batch", "run"): + obj = QuantObjective( + name=f"int8_{scope}", + w_bits=8, + w_group=None, + a_bits=8, + a_mode="per_tensor_static", + a_static_scope=scope, + ) + rs = learn_rotations( + _tiny_llama(), + _calib_batches(), + steps=3, + lr=1.0, + objective_cfg=obj, + seed=0, + log_every=0, + ) + assert max(rs.ortho_audit().values()) < 1e-10 + amax = rs.meta.get("static_act_amax", {}) + assert len(amax) == 7 * N_LAYERS, f"expected {7 * N_LAYERS} static amax entries" + assert all(v > 0 for v in amax.values()) + assert rs.meta["objective"]["a_mode"] == "per_tensor_static" + assert rs.meta["objective"]["a_static_scope"] == scope + assert INT8_DEFAULT_OBJECTIVE.a_static_scope == "batch" # preset default + + +# -------------------------------------------------------------------------------------- +# 3. RotationSet save/load round-trip +# -------------------------------------------------------------------------------------- + + +def test_rotation_set_save_load_roundtrip(): + """save() -> load() reproduces every matrix bitwise (flat fp64 R.bin-format dict), + and load() refuses a non-orthogonal file.""" + rs = learn_rotations( + _tiny_qwen3(), + _calib_batches(), + steps=4, + lr=1.0, + objective_cfg=TINY_W4A4, + seed=0, + log_every=0, + ) + fd, path = tempfile.mkstemp(suffix=".bin") + os.close(fd) + try: + rs.save(path) + rs2 = RotationSet.load(path) + assert set(rs2.rotations) == set(rs.rotations) + for k in rs.rotations: + assert torch.equal(rs2.rotations[k], rs.rotations[k]), f"{k} changed in transit" + # Corrupt one matrix off the manifold: load must refuse. + bad = dict(rs.rotations) + bad["R1"] = bad["R1"] * 1.5 + torch.save(bad, path) + with pytest.raises(ValueError, match="not orthogonal"): + RotationSet.load(path) + # Raw-drift-style file (small non-orthogonal perturbation, like a legacy R.bin + # written without the final retraction): plain load refuses, orthogonalize=True + # retracts and passes. + torch.manual_seed(3) + drifted = dict(rs.rotations) + drifted["R1"] = drifted["R1"] + 3e-4 * torch.randn_like(drifted["R1"]) + torch.save(drifted, path) + try: + RotationSet.load(path) + except ValueError: + pass + else: + raise AssertionError("load() accepted a drifted (non-retracted) R1") + rs3 = RotationSet.load(path, orthogonalize=True) + assert max(rs3.ortho_audit().values()) < 1e-10 + # the retraction stayed near the drifted matrix (order of the perturbation) + assert (rs3.R1 - drifted["R1"]).abs().max().item() < 1e-2 + finally: + os.unlink(path) + + +# -------------------------------------------------------------------------------------- +# 4. fold_rotations external-matrix path +# -------------------------------------------------------------------------------------- + + +def test_fold_external_matches_seed_path(): + """Feeding the seed path's returned matrices back through R1=/R2= reproduces every + parameter bitwise on both architectures — the external path is the same fold.""" + for build in (_tiny_llama, _tiny_qwen3): + model_seed = build() + rots = fold_rotations(model_seed, mode="hadamard", seed=3, use_r2=True) + model_ext = build() # identical construction seed -> identical weights + returned = fold_rotations( + model_ext, R1=rots["R1"], R2={k: v for k, v in rots.items() if k != "R1"} + ) + params_seed = dict(model_seed.named_parameters()) + for n, p in model_ext.named_parameters(): + assert torch.equal(p.data, params_seed[n].data), f"{build.__name__}: {n} differs" + assert set(returned) == set(rots) + for k in rots: + assert torch.equal(returned[k], rots[k]) + + +def test_fold_external_accepts_sequence_r2(): + """R2 as a plain list ordered by layer is equivalent to the keyed-dict form.""" + rots = fold_rotations(_tiny_qwen3(), mode="hadamard", seed=4) + r2_list = [rots[f"model.layers.{i}.self_attn.R2"] for i in range(N_LAYERS)] + model_a, model_b = _tiny_qwen3(), _tiny_qwen3() + fold_rotations(model_a, R1=rots["R1"], R2=r2_list) + fold_rotations(model_b, R1=rots["R1"], R2={k: v for k, v in rots.items() if k != "R1"}) + pb = dict(model_b.named_parameters()) + for n, p in model_a.named_parameters(): + assert torch.equal(p.data, pb[n].data), f"{n} differs between R2 forms" + + +def test_fold_external_validation_errors(): + """Bad external inputs raise ValueError: R2 without R1; use_r2=True with only R1; + non-orthogonal R1; wrong R2 count.""" + rots = fold_rotations(_tiny_qwen3(), mode="hadamard", seed=0) + R1 = rots["R1"] + R2 = {k: v for k, v in rots.items() if k != "R1"} + + def expect_value_error(fn, what): + try: + fn() + except ValueError: + pass + else: + raise AssertionError(f"expected ValueError: {what}") + + expect_value_error(lambda: fold_rotations(_tiny_qwen3(), R2=R2), "R2 without R1") + expect_value_error(lambda: fold_rotations(_tiny_qwen3(), R1=R1), "use_r2=True with only R1") + expect_value_error( + lambda: fold_rotations(_tiny_qwen3(), R1=R1 * 1.5, R2=R2), "non-orthogonal R1" + ) + expect_value_error( + lambda: fold_rotations( + _tiny_qwen3(), R1=R1, R2=[next(iter(R2.values()))] + ), # 1 matrix for 2 layers + "wrong R2 count", + ) + + +def test_fold_external_r1_only_use_r2_false(): + """R1-only external fold (use_r2=False) matches the seed path with use_r2=False.""" + model_seed = _tiny_llama() + rots = fold_rotations(model_seed, mode="hadamard", seed=6, use_r2=False) + assert set(rots) == {"R1"} + model_ext = _tiny_llama() + fold_rotations(model_ext, R1=rots["R1"], use_r2=False) + ps = dict(model_seed.named_parameters()) + for n, p in model_ext.named_parameters(): + assert torch.equal(p.data, ps[n].data), f"{n} differs" + + +if __name__ == "__main__": + tests = [(n, f) for n, f in sorted(globals().items()) if n.startswith("test_") and callable(f)] + failed = [] + for name, fn in tests: + try: + fn() + print(f"PASS {name}", flush=True) + except Exception: + failed.append(name) + print(f"FAIL {name}", flush=True) + traceback.print_exc() + print( + f"\n{len(tests) - len(failed)}/{len(tests)} tests passed" + + (f"; FAILED: {failed}" if failed else "") + ) + sys.exit(1 if failed else 0) diff --git a/tests/unit/torch/quantization/test_rotation_paper_objective.py b/tests/unit/torch/quantization/test_rotation_paper_objective.py new file mode 100644 index 00000000000..04d578c69b8 --- /dev/null +++ b/tests/unit/torch/quantization/test_rotation_paper_objective.py @@ -0,0 +1,250 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the paper-protocol objective extensions: per-token ASYM min-max +activation fake-quant (``QuantObjective.a_asym``) and the training-graph-only online R4 +down_proj Hadamard (``QuantObjective.r4_in_graph``). + +Plain test_* functions with asserts: collectable by pytest, and also runnable without it +via ``python test_rotation_paper_objective.py``. CPU-only, tiny models, seconds per test. +""" + +import sys +import traceback + +import pytest +import torch +from transformers import LlamaConfig, LlamaForCausalLM + +from modelopt.torch.quantization.rotation import ( + W16A4_ASYM_R4G_OBJECTIVE, + QuantObjective, + learn_rotations, +) +from modelopt.torch.quantization.rotation.learn import _fq_act, _fq_act_asym, _walsh_hadamard + +VOCAB = 128 +HIDDEN = 64 +HEAD_DIM = 32 +N_LAYERS = 2 + + +def _randomize_rmsnorm_gains(model): + for module in model.modules(): + if type(module).__name__.endswith("RMSNorm"): + module.weight.data = 1.0 + 0.1 * torch.randn_like(module.weight.data) + + +def _tiny_llama(intermediate=2 * HIDDEN): + torch.manual_seed(1234) + cfg = LlamaConfig( + vocab_size=VOCAB, + hidden_size=HIDDEN, + intermediate_size=intermediate, # default 128 = 2^7: power-of-2 R4 seam + num_hidden_layers=N_LAYERS, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=HEAD_DIM, + max_position_embeddings=128, + tie_word_embeddings=False, + attn_implementation="eager", + ) + model = LlamaForCausalLM(cfg).eval() + _randomize_rmsnorm_gains(model) + return model + + +def _calib_batches(n_batches=2, bs=2, seq=16, seed=7): + torch.manual_seed(seed) + return [torch.randint(0, VOCAB, (bs, seq)) for _ in range(n_batches)] + + +# -------------------------------------------------------------------------------------- +# 1. Asym activation fake-quant: official ActQuantizer numerics, exactly +# -------------------------------------------------------------------------------------- + + +def _official_asym_reference(x: torch.Tensor, bits: int) -> torch.Tensor: + """Line-for-line transcription of the official SpinQuant ActQuantizer + (utils/quant_utils.py) at sym=False, clip_ratio=1, groupsize=-1: find_params' + zero-inclusive per-token range + asym_quant_dequant.""" + maxq = float(2**bits - 1) + flat = x.reshape(-1, x.shape[-1]) + tmp = torch.zeros(flat.shape[0], dtype=x.dtype) + xmin = torch.minimum(flat.min(1)[0], tmp) + xmax = torch.maximum(flat.max(1)[0], tmp) + degen = (xmin == 0) & (xmax == 0) + xmin[degen] = -1 + xmax[degen] = +1 + scale = ((xmax - xmin) / maxq).unsqueeze(1) + zero = torch.round(-xmin.unsqueeze(1) / scale) + q = torch.clamp(torch.round(flat / scale) + zero, 0, maxq) + return (scale * (q - zero)).reshape(x.shape) + + +def test_asym_act_quant_matches_official_reference(): + """_fq_act_asym reproduces the official asym recipe bit-for-bit on hard cases: + mixed-sign tokens, an ALL-POSITIVE token (zero-inclusion changes the range), an + all-negative token, and an all-zero token (the [-1, 1] degenerate fallback).""" + torch.manual_seed(0) + x = torch.randn(5, 4, 32) + x[0, 0] = x[0, 0].abs() + 0.5 # all-positive token: xmin must clamp to 0 + x[1, 1] = -x[1, 1].abs() - 0.5 # all-negative token: xmax must clamp to 0 + x[2, 2] = 0.0 # degenerate all-zero token + for bits in (4, 8): + got = _fq_act_asym(x, bits) + want = _official_asym_reference(x, bits) + assert torch.equal(got, want), ( + f"bits={bits}: asym fake-quant deviates from the official recipe " + f"(max |diff| = {(got - want).abs().max().item():.3e})" + ) + + +def test_asym_ste_gradient_is_identity(): + torch.manual_seed(1) + x = torch.randn(3, 8, 16, requires_grad=True) + _fq_act_asym(x, 4).sum().backward() + assert torch.equal(x.grad, torch.ones_like(x)), "STE gradient must be identity" + + +def test_asym_beats_sym_on_shifted_activations(): + """The reason the paper uses asym (A.4): a positively-shifted activation (post-SiLU + regime) wastes half the sym grid. Same tensor, same bits: asym error must be well + below sym error.""" + torch.manual_seed(2) + x = torch.randn(4, 16, 64) + 3.0 # strong positive shift + bits = 4 + qpos = float(2 ** (bits - 1) - 1) + s = (x.abs().amax(dim=-1, keepdim=True) / qpos).clamp_min(1e-12) + err_sym = (_fq_act(x, s, bits) - x).norm() + err_asym = (_fq_act_asym(x, bits) - x).norm() + assert err_asym < 0.7 * err_sym, ( + f"asym ({err_asym:.4f}) should clearly beat sym ({err_sym:.4f}) on shifted data" + ) + + +# -------------------------------------------------------------------------------------- +# 2. Walsh-Hadamard helper +# -------------------------------------------------------------------------------------- + + +def test_walsh_hadamard_properties(): + for n in (1, 2, 8, 128): + H = _walsh_hadamard(n) + assert H.shape == (n, n) and H.dtype == torch.float32 + assert torch.equal(H, H.t()), "Sylvester Hadamard must be symmetric" + err = (H @ H.t() - torch.eye(n)).abs().max().item() + assert err < 1e-6, f"n={n}: |H H^T - I| = {err:.3e}" + for bad in (0, 3, 48, 6144): # 6144 = Qwen3-1.7B intermediate — documented unsupported + try: + _walsh_hadamard(bad) + raise AssertionError(f"n={bad} should have raised NotImplementedError") + except NotImplementedError: + pass + + +# -------------------------------------------------------------------------------------- +# 3. r4_in_graph: functional identity in the graph, absent from the output +# -------------------------------------------------------------------------------------- + + +def test_r4_in_graph_is_functional_identity(): + """With ALL quantizers off, the r4 pair (input hook x @ H + weight cols @ H) must be + a functional identity: at lr=0 the step-0 loss equals the objective=None loss, and + the returned rotations equal the objective=None run's rotations (same seed draws + + final retraction) — i.e. no H leaks into the deployable output.""" + batches = _calib_batches() + r4_only = QuantObjective( + name="r4_only", w_bits=None, w_group=None, a_bits=None, r4_in_graph=True + ) + rs_r4 = learn_rotations( + _tiny_llama(), batches, steps=1, lr=0.0, objective_cfg=r4_only, seed=3, log_every=0 + ) + rs_ref = learn_rotations( + _tiny_llama(), batches, steps=1, lr=0.0, objective_cfg=None, seed=3, log_every=0 + ) + l_r4, l_ref = rs_r4.history[0]["loss"], rs_ref.history[0]["loss"] + assert abs(l_r4 - l_ref) < 1e-3 * max(1.0, abs(l_ref)), ( + f"r4-only step-0 loss {l_r4} != objective-None loss {l_ref} — the H pair is not " + "a functional identity" + ) + for k in rs_ref.rotations: + d = (rs_r4.rotations[k] - rs_ref.rotations[k]).abs().max().item() + assert d < 1e-5, f"{k}: rotations differ by {d:.3e} — H leaked into the output" + + +def test_r4_rejects_non_pow2_seam(): + model = _tiny_llama(intermediate=96) # 96 = 3 * 32: not a power of 2 + with pytest.raises(NotImplementedError, match="power of 2"): + learn_rotations( + model, + _calib_batches(), + steps=1, + lr=0.0, + objective_cfg=QuantObjective(name="r4", w_bits=None, a_bits=None, r4_in_graph=True), + seed=0, + log_every=0, + ) + + +# -------------------------------------------------------------------------------------- +# 4. The paper preset trains: loss decreases on a tiny overfit case, meta records flags +# -------------------------------------------------------------------------------------- + + +def test_w16a4_asym_r4g_preset_trains_and_records_meta(): + torch.manual_seed(11) + batch = [torch.randint(0, VOCAB, (2, 16))] # single repeated batch: overfit regime + rs = learn_rotations( + _tiny_llama(), + batch, + steps=12, + lr=0.5, + objective_cfg=W16A4_ASYM_R4G_OBJECTIVE, + seed=5, + log_every=0, + ) + first, last = rs.history[0]["loss"], rs.history[-1]["loss"] + assert last < first, f"loss did not decrease: {first} -> {last}" + obj = rs.meta["objective"] + assert obj["a_asym"] is True and obj["r4_in_graph"] is True and obj["w_bits"] is None + audit = rs.ortho_audit() + assert max(audit.values()) < 1e-4, f"rotations left the manifold: {audit}" + + +def test_validation_errors(): + for kwargs in ( + {"a_asym": True, "a_bits": None}, # asym needs a_bits + {"a_asym": True, "a_bits": 8, "a_mode": "per_tensor_static"}, # asym is per-token + ): + try: + QuantObjective(name="bad", w_bits=None, w_group=None, **kwargs) + raise AssertionError(f"{kwargs} should have raised ValueError") + except ValueError: + pass + + +if __name__ == "__main__": + failures = 0 + for name, fn in sorted(globals().items()): + if name.startswith("test_") and callable(fn): + try: + fn() + print(f"PASS {name}") + except Exception: + failures += 1 + print(f"FAIL {name}") + traceback.print_exc() + sys.exit(1 if failures else 0) diff --git a/tests/unit/torch/quantization/test_rotation_transform_qat.py b/tests/unit/torch/quantization/test_rotation_transform_qat.py new file mode 100644 index 00000000000..613bbd64790 --- /dev/null +++ b/tests/unit/torch/quantization/test_rotation_transform_qat.py @@ -0,0 +1,708 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Transform-QAT tests: jointly learned rotations R1/R2 + per-input-channel seam +diagonals (OSTQuant-style) in modelopt.torch.quantization.rotation. + +Covers: + 1. Default-off backward compatibility: ``learn_seam_diag=False`` (present or absent) + is bitwise identical — matching rotations/histories, ``seam_diags is None``, and + the no-diag assembly path reproduces the pre-change formula bitwise. + 2. learn_seam_diag=True: gradients reach R1, every R2 AND every seam-diag parameter + (and nothing else); loss decreases on a repeated-batch overfit; the assembled + reparametrized model is function-preserving at the zeros init AND for arbitrary + nonzero log-scales (quant off) — the structural seam-identity check. + 3. fold_seam_diags round-trip: learn 5 steps -> bake diags + R into a fresh model + (fold_seam_diags then fold_rotations) -> logits match the assembled reparametrized + model, and the folded weights match the assembly entry-for-entry; either fold + order preserves the function. + 4. save/load round-trip including seam_diags; old-format (flat R.bin) files still + load with seam_diags=None; seam_diags=None saves the legacy flat format. + +Plain test_* functions with asserts: collectable by pytest, and also runnable without +it via ``python test_rotation_transform_qat.py`` (the __main__ driver runs every test +function and exits nonzero on any failure). CPU-only, tiny models, seconds per test. +""" + +import os + +os.environ["CUDA_VISIBLE_DEVICES"] = "" # HARD CONSTRAINT: CPU-only, never touch the GPU + +import sys +import tempfile +import traceback + +import pytest +import torch +import torch.nn as nn +from torch.nn.utils import stateless +from transformers import LlamaConfig, LlamaForCausalLM, Qwen3Config, Qwen3ForCausalLM + +from modelopt.torch.quantization.rotation import ( + QuantObjective, + RotationSet, + fold_rotations, + fold_seam_diags, + learn_rotations, +) +from modelopt.torch.quantization.rotation.learn import ( + _ATTN_PROJS, + _MLP_PROJS, + _SEAM_DIAGS_KEY, + _ActQuantHooks, + _assemble_effective_weights, + _fq_weight, +) + +VOCAB = 128 +# Standard tiny fixture (same as test_rotation_{fold,learn,ext_learner}.py): HEAD_DIM +# deliberately decoupled from HIDDEN // num_attention_heads (32 vs 64//4 = 16), like +# Qwen3-0.6B. GQA is real (4 q heads on 2 kv heads), so the o-seam group expansion is +# exercised, not degenerate. +HIDDEN = 64 +HEAD_DIM = 32 +N_LAYERS = 2 +N_KV = 2 +INTERMEDIATE = 2 * HIDDEN # 128 +O_SEAM_DIM = N_KV * HEAD_DIM # 64 + +# Tiny-model W4A4 (g=16 divides every in_features {64, 128, 128}), as in the sibling +# test files — once plain, once with the transform-QAT flag. +TINY_W4A4 = QuantObjective( + name="tiny_w4a4", w_bits=4, w_group=16, a_bits=4, a_mode="per_token_dynamic" +) +TINY_W4A4_DIAG = QuantObjective( + name="tiny_w4a4_diag", + w_bits=4, + w_group=16, + a_bits=4, + a_mode="per_token_dynamic", + learn_seam_diag=True, +) + + +def _randomize_rmsnorm_gains(model): + for module in model.modules(): + if type(module).__name__.endswith("RMSNorm"): + module.weight.data = 1.0 + 0.1 * torch.randn_like(module.weight.data) + + +def _build(cfg_cls, model_cls, tie=False): + torch.manual_seed(1234) + cfg = cfg_cls( + vocab_size=VOCAB, + hidden_size=HIDDEN, + intermediate_size=INTERMEDIATE, + num_hidden_layers=N_LAYERS, + num_attention_heads=4, + num_key_value_heads=N_KV, + head_dim=HEAD_DIM, + max_position_embeddings=128, + tie_word_embeddings=tie, + attn_implementation="eager", + ) + model = model_cls(cfg).eval() + _randomize_rmsnorm_gains(model) + return model + + +def _tiny_llama(tie=False): + return _build(LlamaConfig, LlamaForCausalLM, tie) + + +def _tiny_qwen3(tie=False): + return _build(Qwen3Config, Qwen3ForCausalLM, tie) + + +def _calib_batches(n_batches=2, bs=2, seq=16, seed=7): + torch.manual_seed(seed) + return [torch.randint(0, VOCAB, (bs, seq)) for _ in range(n_batches)] + + +def _logits(model): + torch.manual_seed(99) + ids = torch.randint(0, VOCAB, (2, 8)) + with torch.no_grad(): + return model(ids).logits + + +def _base_weights(model): + sd = dict(model.named_parameters()) + base = { + "model.embed_tokens.weight": sd["model.embed_tokens.weight"].data, + "lm_head.weight": sd["lm_head.weight"].data, + } + for i in range(len(model.model.layers)): + for proj in _ATTN_PROJS + _MLP_PROJS: + name = f"model.layers.{i}.{proj}.weight" + base[name] = sd[name].data + return base + + +def _r_leaves(rotations, n_layers): + R1 = nn.Parameter(rotations["R1"].to(torch.float32)) + R2s = [ + nn.Parameter(rotations[f"model.layers.{i}.self_attn.R2"].to(torch.float32)) + for i in range(n_layers) + ] + return R1, R2s + + +def _diag_leaves(n_layers, fill=None): + """Per-layer {down, o} log-scale leaves: zeros (identity) or a provided filler fn.""" + out = [] + for i in range(n_layers): + if fill is None: + d = torch.zeros(INTERMEDIATE) + o = torch.zeros(O_SEAM_DIM) + else: + d, o = fill(i) + out.append({"down": nn.Parameter(d.float()), "o": nn.Parameter(o.float())}) + return out + + +def _assembled_logits(model, eff): + torch.manual_seed(99) + ids = torch.randint(0, VOCAB, (2, 8)) + with torch.no_grad(), stateless._reparametrize_module(model, eff): + return model(ids).logits + + +# -------------------------------------------------------------------------------------- +# 1. Default-off = bitwise identical behavior to before +# -------------------------------------------------------------------------------------- + + +def _assemble_pre_change(base, R1, R2s, n_layers, head_dim, objective, out_dtype): + """Rotation-only reference assembly (no seam diagonals) — the bitwise + oracle for the learn_seam_diag=False path.""" + compute = R1.dtype + d = head_dim + + def fin(w): + if objective is not None and objective.w_bits is not None: + w = _fq_weight(w, objective) + return w.to(out_dtype) + + eff = {} + for name in ("model.embed_tokens.weight", "lm_head.weight"): + eff[name] = (base[name].to(compute) @ R1).to(out_dtype) + + for i in range(n_layers): + R2 = R2s[i] + pre = f"model.layers.{i}." + for proj in ("self_attn.q_proj", "self_attn.k_proj", "mlp.gate_proj", "mlp.up_proj"): + n = pre + proj + ".weight" + eff[n] = fin(base[n].to(compute) @ R1) + + n = pre + "mlp.down_proj.weight" + eff[n] = fin(R1.t() @ base[n].to(compute)) + + n = pre + "self_attn.v_proj.weight" + a = base[n].to(compute) @ R1 + o_f, i_f = a.shape + a = (a.t().reshape(i_f, o_f // d, d) @ R2).reshape(i_f, o_f).t().contiguous() + eff[n] = fin(a) + + n = pre + "self_attn.o_proj.weight" + w = R1.t() @ base[n].to(compute) + o_f, i_f = w.shape + eff[n] = fin((w.reshape(o_f, i_f // d, d) @ R2).reshape(o_f, i_f)) + return eff + + +def test_default_off_bitwise_identical(): + """learn_seam_diag defaults to False, and False (explicit or absent) is one code + path: (i) steps=3 runs with the field absent vs. explicitly False give bitwise-equal + rotations and identical loss histories, both with seam_diags=None; (ii) the no-diag + assembly reproduces the pre-change formula bitwise for every effective weight.""" + assert QuantObjective(name="x").learn_seam_diag is False # default off + explicit_off = QuantObjective( + name="tiny_w4a4", + w_bits=4, + w_group=16, + a_bits=4, + a_mode="per_token_dynamic", + learn_seam_diag=False, + ) + assert explicit_off == TINY_W4A4 # frozen dataclass equality: absent == False + + runs = [] + for obj in (TINY_W4A4, explicit_off): + rs = learn_rotations( + _tiny_qwen3(), + _calib_batches(), + steps=3, + lr=1.0, + objective_cfg=obj, + seed=0, + log_every=0, + ) + assert rs.seam_diags is None, "learn_seam_diag=False must not produce seam_diags" + assert "seam_diag" not in rs.meta + runs.append(rs) + a, b = runs + assert set(a.rotations) == set(b.rotations) + for k in a.rotations: + assert torch.equal(a.rotations[k], b.rotations[k]), f"{k}: off-path not bitwise" + assert [r["loss"] for r in a.history] == [r["loss"] for r in b.history] + + # Assembly-level bitwise oracle vs. the verbatim pre-change formula. + model = _tiny_qwen3() + init = learn_rotations(model, [], steps=0, objective_cfg=None, seed=0, log_every=0) + R1, R2s = _r_leaves(init.rotations, N_LAYERS) + base = _base_weights(model) + for obj in (TINY_W4A4, None): + new = _assemble_effective_weights( + base, + R1, + R2s, + N_LAYERS, + HEAD_DIM, + obj, + torch.float32, + seam_diag_params=None, + ) + old = _assemble_pre_change(base, R1, R2s, N_LAYERS, HEAD_DIM, obj, torch.float32) + assert set(new) == set(old) + for k in old: + assert torch.equal(new[k], old[k]), f"{k}: no-diag assembly != pre-change (obj={obj})" + + +def test_step0_loss_identical_diag_on_vs_off(): + """At the zeros init the diagonals are exact identity scales, so the step-0 loss of a + learn_seam_diag=True run equals the rotation-only run's bitwise (trajectories may + diverge from step 1 once Adam moves the diagonals).""" + on = learn_rotations( + _tiny_qwen3(), + _calib_batches(), + steps=2, + lr=1.0, + objective_cfg=TINY_W4A4_DIAG, + seed=0, + log_every=0, + ) + off = learn_rotations( + _tiny_qwen3(), + _calib_batches(), + steps=2, + lr=1.0, + objective_cfg=TINY_W4A4, + seed=0, + log_every=0, + ) + assert on.history[0]["loss"] == off.history[0]["loss"], ( + f"step-0 loss differs: diag-on {on.history[0]['loss']} vs off {off.history[0]['loss']}" + ) + + +# -------------------------------------------------------------------------------------- +# 2. learn_seam_diag=True: gradients, descent, structural function preservation +# -------------------------------------------------------------------------------------- + + +def test_grads_reach_rotations_and_diags_and_nothing_else(): + """One step-0 forward/backward with diag leaves (zeros init) on both archs: finite + loss; non-None, finite, nonzero grads on R1, every R2, and every down/o log-scale; + no grad on any model parameter.""" + for build in (_tiny_llama, _tiny_qwen3): + model = build() + init = learn_rotations(model, [], steps=0, objective_cfg=None, seed=0, log_every=0) + R1, R2s = _r_leaves(init.rotations, N_LAYERS) + diag = _diag_leaves(N_LAYERS) + base = _base_weights(model) + hooks = _ActQuantHooks(TINY_W4A4_DIAG) + n_hooked = hooks.attach(model) + assert n_hooked == 7 * N_LAYERS + torch.manual_seed(42) + ids = torch.randint(0, VOCAB, (2, 16)) + try: + eff = _assemble_effective_weights( + base, + R1, + R2s, + N_LAYERS, + HEAD_DIM, + TINY_W4A4_DIAG, + torch.float32, + seam_diag_params=diag, + ) + with stateless._reparametrize_module(model, eff): + loss = model(input_ids=ids, labels=ids, use_cache=False).loss + loss.backward() + finally: + hooks.remove() + tag = build.__name__ + assert torch.isfinite(loss), f"{tag}: loss not finite" + assert R1.grad is not None and R1.grad.abs().max() > 0, f"{tag}: R1 grad missing/zero" + for i, r2 in enumerate(R2s): + assert r2.grad is not None and torch.isfinite(r2.grad).all(), f"{tag}: R2[{i}]" + assert r2.grad.abs().max() > 0, f"{tag}: R2[{i}] grad identically zero" + for i, sp in enumerate(diag): + for key in ("down", "o"): + g = sp[key].grad + assert g is not None, f"{tag}: no grad reached log_s_{key}[{i}]" + assert torch.isfinite(g).all(), f"{tag}: log_s_{key}[{i}] grad not finite" + assert g.abs().max() > 0, f"{tag}: log_s_{key}[{i}] grad identically zero" + for name, p in model.named_parameters(): + assert p.grad is None, f"{tag}: model param {name} received a gradient" + + +def test_learn_with_diag_trains_and_moves_scales(): + """learn_rotations with learn_seam_diag=True: seam_diags has every layer with the + right shapes, strictly positive, visibly moved off the identity; the rotations stay + orthonormal; meta records the diag group; loss decreases on a repeated batch.""" + batch = _calib_batches(n_batches=1, bs=2, seq=32, seed=11) + rs = learn_rotations( + _tiny_qwen3(), + batch, + steps=25, + lr=1.0, + objective_cfg=TINY_W4A4_DIAG, + seed=0, + log_every=0, + ) + assert rs.seam_diags is not None and set(rs.seam_diags) == set(range(N_LAYERS)) + moved = 0.0 + for i in range(N_LAYERS): + sd, so = rs.seam_diags[i]["down"], rs.seam_diags[i]["o"] + assert sd.shape == (INTERMEDIATE,) and so.shape == (O_SEAM_DIM,) + assert sd.dtype == torch.float64 and so.dtype == torch.float64 + assert sd.device.type == "cpu" and so.device.type == "cpu" + assert (sd > 0).all() and (so > 0).all() + moved = max(moved, (sd - 1).abs().max().item(), (so - 1).abs().max().item()) + assert moved > 1e-4, f"seam scales never moved off identity (max |s-1| = {moved:.3e})" + assert max(rs.ortho_audit().values()) < 1e-10 + assert rs.meta["objective"]["learn_seam_diag"] is True + assert rs.meta["seam_diag"]["lr"] == 1e-2 + assert 0 < rs.meta["seam_diag"]["s_min"] <= rs.meta["seam_diag"]["s_max"] + losses = [r["loss"] for r in rs.history] + assert min(losses[-5:]) < losses[0], ( + f"no decrease with diag learning: first {losses[0]:.6f}, last5 {losses[-5:]}" + ) + + +def test_steps0_diag_identity_and_rng_stream_unchanged(): + """steps=0 with learn_seam_diag=True returns identity scales (exact ones — zeros + init consumes no RNG) and the SAME bitwise rotation draws as the seed path — the + diag machinery must not shift the RNG stream.""" + init = learn_rotations( + _tiny_qwen3(), [], steps=0, objective_cfg=TINY_W4A4_DIAG, seed=5, log_every=0 + ) + assert init.seam_diags is not None + for i in range(N_LAYERS): + for key, dim in (("down", INTERMEDIATE), ("o", O_SEAM_DIM)): + s = init.seam_diags[i][key] + assert torch.equal(s, torch.ones(dim, dtype=torch.float64)), ( + f"steps=0 seam scale [{i}][{key}] is not exact identity" + ) + folded = fold_rotations(_tiny_qwen3(), mode="hadamard", seed=5, use_r2=True) + for k in folded: + assert torch.equal(init.rotations[k], folded[k]), f"{k}: RNG stream shifted" + + +def test_assembled_function_preserving_at_init_and_any_diag(): + """Quant off: the assembled reparametrized model equals the plain model function + (logit-level) BOTH at the zeros init and — the real structural check — for random + nonzero log-scales at every seam (any positive diagonal is an exact identity; a + wrong axis, wrong GQA expansion, or wrong diag/R2 order would be O(1) off).""" + for build in (_tiny_llama, _tiny_qwen3): + plain = _logits(build()) + model = build() + init = learn_rotations(model, [], steps=0, objective_cfg=None, seed=0, log_every=0) + R1, R2s = _r_leaves(init.rotations, N_LAYERS) + base = _base_weights(model) + + def rand_fill(i): + g = torch.Generator().manual_seed(100 + i) + return ( + 0.4 * torch.randn(INTERMEDIATE, generator=g), + 0.4 * torch.randn(O_SEAM_DIM, generator=g), + ) + + for label, fill in (("zeros-init", None), ("random-diag", rand_fill)): + diag = _diag_leaves(N_LAYERS, fill=fill) + if label == "random-diag": # make sure the case is not vacuous + assert max(p.abs().max().item() for sp in diag for p in sp.values()) > 0.3 + eff = _assemble_effective_weights( + base, + R1, + R2s, + N_LAYERS, + HEAD_DIM, + None, + torch.float32, + seam_diag_params=diag, + ) + out = _assembled_logits(model, eff) + dmax = (out - plain).abs().max().item() + assert torch.allclose(out, plain, rtol=0, atol=1e-3), ( + f"{build.__name__}/{label}: assembled model broke function " + f"preservation (max |delta logit| = {dmax:.3e})" + ) + + +# -------------------------------------------------------------------------------------- +# 3. fold_seam_diags round-trip against the assembled reparametrized model +# -------------------------------------------------------------------------------------- + + +def test_fold_seam_diags_roundtrip_matches_assembly(): + """Learn 5 steps (diag on) -> bake into a fresh model with fold_seam_diags then + fold_rotations -> (i) logits match the assembled reparametrized model within bf16- + level tolerance, (ii) the folded weights match the assembly's effective weights + entry-for-entry (the two folds compose to exactly the assembly's prefold-inside / + rotation-outside convention), (iii) the reverse fold order preserves the function + too, and (iv) the fp16-safety clamp never bit.""" + rs = learn_rotations( + _tiny_qwen3(), + _calib_batches(), + steps=5, + lr=1.0, + objective_cfg=TINY_W4A4_DIAG, + seed=0, + log_every=0, + ) + moved = max( + (rs.seam_diags[i][k] - 1).abs().max().item() for i in range(N_LAYERS) for k in ("down", "o") + ) + assert moved > 1e-4, "trained scales are still identity — round-trip would be vacuous" + + # Assembled reference: prepped (untied+fused) model reparametrized with the learned + # R and log-scales, quantization off. + model_asm = _tiny_qwen3() + learn_rotations(model_asm, [], steps=0, objective_cfg=None, seed=0, log_every=0) + R1, R2s = _r_leaves(rs.rotations, N_LAYERS) + diag = _diag_leaves( + N_LAYERS, + fill=lambda i: ( + torch.log(rs.seam_diags[i]["down"]).float(), + torch.log(rs.seam_diags[i]["o"]).float(), + ), + ) + eff = _assemble_effective_weights( + _base_weights(model_asm), + R1, + R2s, + N_LAYERS, + HEAD_DIM, + None, + torch.float32, + seam_diag_params=diag, + ) + ref = _assembled_logits(model_asm, eff) + + # Fold path on a fresh model: diags first, then rotations (assembly convention). + model_fold = _tiny_qwen3() + evidence = fold_seam_diags(model_fold, rs.seam_diags) + assert set(evidence["layers"]) == set(range(N_LAYERS)) + assert not any(rec["clamped"] for rec in evidence["layers"].values()) + assert all(rec["down_s_spread"] > 1 for rec in evidence["layers"].values()) + fold_rotations(model_fold, R1=rs.R1, R2=rs.R2) + out = _logits(model_fold) + dmax = (out - ref).abs().max().item() + assert torch.allclose(out, ref, rtol=0, atol=1e-2), ( + f"folded model != assembled model (max |delta logit| = {dmax:.3e})" + ) + + # Weight-level: every effective weight is reproduced by the composed folds. + params = dict(model_fold.named_parameters()) + for name, w in eff.items(): + d = (params[name].data - w).abs().max().item() + assert d < 1e-4, f"{name}: folded weight != assembled effective weight ({d:.3e})" + + # Reverse order (rotations first, then diags) is also a functional identity. + model_rev = _tiny_qwen3() + fold_rotations(model_rev, R1=rs.R1, R2=rs.R2) + fold_seam_diags(model_rev, rs.seam_diags) + out_rev = _logits(model_rev) + assert torch.allclose(out_rev, ref, rtol=0, atol=1e-2), ( + f"reverse fold order broke function (max delta {(out_rev - ref).abs().max().item():.3e})" + ) + + +def test_fold_seam_diags_validation_and_clamp(): + """fold_seam_diags refuses wrong-length or non-positive scales and unknown layers; + the smax clamp bites (and is reported) for scales beyond the ceiling; a plain + identity fold is a bf16-free exact no-op on an fp32 model up to fp64 round-trip.""" + ones_d, ones_o = torch.ones(INTERMEDIATE), torch.ones(O_SEAM_DIM) + + def expect_error(fn, what, exc=ValueError): + try: + fn() + except exc: + pass + else: + raise AssertionError(f"expected {exc.__name__}: {what}") + + expect_error( + lambda: fold_seam_diags(_tiny_qwen3(), {0: {"down": torch.ones(7), "o": ones_o}}), + "wrong down length", + ) + expect_error( + lambda: fold_seam_diags(_tiny_qwen3(), {0: {"down": -ones_d, "o": ones_o}}), + "negative scales", + ) + expect_error( + lambda: fold_seam_diags(_tiny_qwen3(), {99: {"down": ones_d, "o": ones_o}}), + "layer index out of range", + ) + expect_error( + lambda: fold_seam_diags(_tiny_qwen3(), {0: {"down": ones_d}}), + "missing 'o' key", + ) + + # Identity scales: exact no-op (fp64 round-trip of unchanged values is bitwise). + model = _tiny_qwen3() + before = {n: p.data.clone() for n, p in model.named_parameters()} + ev = fold_seam_diags(model, {i: {"down": ones_d, "o": ones_o} for i in range(N_LAYERS)}) + for n, p in model.named_parameters(): + assert torch.equal(p.data, before[n]), f"identity fold changed {n}" + assert not any(rec["clamped"] for rec in ev["layers"].values()) + + # Clamp: s=512 with smax=256 folds as 256 and reports clamped=True. + model = _tiny_qwen3() + up0 = model.model.layers[0].mlp.up_proj.weight.data.clone() + ev = fold_seam_diags(model, {0: {"down": 512.0 * ones_d, "o": ones_o}}, smax=256.0) + assert ev["layers"][0]["clamped"] is True + assert ev["layers"][0]["down_s_max"] == 512.0 # telemetry reports the raw scale + got = model.model.layers[0].mlp.up_proj.weight.data + want = (up0.to(torch.float64) / 256.0).to(torch.float32) + assert torch.equal(got, want), "clamped fold did not use the smax ceiling" + + +# -------------------------------------------------------------------------------------- +# 4. save/load round-trip incl. seam_diags; old-format compatibility +# -------------------------------------------------------------------------------------- + + +def test_save_load_roundtrip_with_seam_diags_and_old_format(): + """New format round-trips rotations AND seam_diags bitwise; seam_diags=None writes + the legacy flat dict (no reserved key); an old-format (pure-rotation) file loads + with seam_diags=None; the ortho gate and orthogonalize=True still work with the + seam payload present.""" + rs = learn_rotations( + _tiny_qwen3(), + _calib_batches(), + steps=4, + lr=1.0, + objective_cfg=TINY_W4A4_DIAG, + seed=0, + log_every=0, + ) + assert rs.seam_diags is not None + fd, path = tempfile.mkstemp(suffix=".bin") + os.close(fd) + try: + # New-format round-trip: everything bitwise. + rs.save(path) + raw = torch.load(path, map_location="cpu", weights_only=True) + assert _SEAM_DIAGS_KEY in raw + rs2 = RotationSet.load(path) + assert set(rs2.rotations) == set(rs.rotations) + for k in rs.rotations: + assert torch.equal(rs2.rotations[k], rs.rotations[k]), f"{k} changed in transit" + assert rs2.seam_diags is not None and set(rs2.seam_diags) == set(rs.seam_diags) + for i in rs.seam_diags: + for key in ("down", "o"): + assert torch.equal(rs2.seam_diags[i][key], rs.seam_diags[i][key]), ( + f"seam_diags[{i}][{key}] changed in transit" + ) + + # seam_diags=None saves the legacy flat format: only rotation keys on disk. + rs_plain = RotationSet(rotations=dict(rs.rotations)) + rs_plain.save(path) + raw = torch.load(path, map_location="cpu", weights_only=True) + assert set(raw) == set(rs.rotations), "legacy save format changed" + + # Old-format file (flat rotation dict, e.g. a legacy R.bin): loads fine, + # seam_diags=None. + torch.save(dict(rs.rotations), path) + rs3 = RotationSet.load(path) + assert rs3.seam_diags is None + for k in rs.rotations: + assert torch.equal(rs3.rotations[k], rs.rotations[k]) + + # Ortho gate still guards new-format files. + bad = dict(rs.rotations) + bad["R1"] = bad["R1"] * 1.5 + bad[_SEAM_DIAGS_KEY] = { + int(i): {k: v.clone() for k, v in pair.items()} for i, pair in rs.seam_diags.items() + } + torch.save(bad, path) + with pytest.raises(ValueError, match="not orthogonal"): + RotationSet.load(path) + + # orthogonalize=True retracts the rotations and PRESERVES the seam payload. + torch.manual_seed(3) + drifted = dict(rs.rotations) + drifted["R1"] = drifted["R1"] + 3e-4 * torch.randn_like(drifted["R1"]) + drifted[_SEAM_DIAGS_KEY] = bad[_SEAM_DIAGS_KEY] + torch.save(drifted, path) + rs4 = RotationSet.load(path, orthogonalize=True) + assert max(rs4.ortho_audit().values()) < 1e-10 + assert rs4.seam_diags is not None + for i in rs.seam_diags: + for key in ("down", "o"): + assert torch.equal(rs4.seam_diags[i][key], rs.seam_diags[i][key]) + finally: + os.unlink(path) + + +def test_rotation_set_rejects_bad_seam_diags(): + """The RotationSet constructor validates the seam payload: non-positive scales and + wrong key sets are refused; int-like string layer keys are normalized to int.""" + rots = fold_rotations(_tiny_qwen3(), mode="hadamard", seed=0) + good = { + i: {"down": torch.ones(INTERMEDIATE), "o": torch.ones(O_SEAM_DIM)} for i in range(N_LAYERS) + } + rs = RotationSet(rotations=dict(rots), seam_diags={str(i): v for i, v in good.items()}) + assert set(rs.seam_diags) == set(range(N_LAYERS)) # str keys normalized + assert rs.seam_diags[0]["down"].dtype == torch.float64 + + def expect_value_error(seam, what): + try: + RotationSet(rotations=dict(rots), seam_diags=seam) + except ValueError: + pass + else: + raise AssertionError(f"expected ValueError: {what}") + + neg = {0: {"down": -torch.ones(INTERMEDIATE), "o": torch.ones(O_SEAM_DIM)}} + expect_value_error(neg, "negative scales") + zero = {0: {"down": torch.zeros(INTERMEDIATE), "o": torch.ones(O_SEAM_DIM)}} + expect_value_error(zero, "zero scales") + missing = {0: {"down": torch.ones(INTERMEDIATE)}} + expect_value_error(missing, "missing 'o' key") + extra = {0: {"down": torch.ones(INTERMEDIATE), "o": torch.ones(O_SEAM_DIM), "x": 1}} + expect_value_error(extra, "extra key") + + +if __name__ == "__main__": + tests = [(n, f) for n, f in sorted(globals().items()) if n.startswith("test_") and callable(f)] + failed = [] + for name, fn in tests: + try: + fn() + print(f"PASS {name}", flush=True) + except Exception: + failed.append(name) + print(f"FAIL {name}", flush=True) + traceback.print_exc() + print( + f"\n{len(tests) - len(failed)}/{len(tests)} tests passed" + + (f"; FAILED: {failed}" if failed else "") + ) + sys.exit(1 if failed else 0)