From bbb3630760e92edf211caa4655815d8f8179b0f8 Mon Sep 17 00:00:00 2001 From: Jie Ren Date: Mon, 3 Aug 2026 15:15:41 -0700 Subject: [PATCH 1/5] Add offline SpinQuant/QuaRot rotation folding (R1+R2) as a pre-quantization checkpoint transform fold_rotations folds a global residual-stream rotation R1 and per-layer head-space rotations R2 (v_proj -> o_proj path, GQA-exact, one shared R2 per layer) into HF RMSNorm decoder weights in place (Llama family, Qwen3), after fusing RMSNorm gains into downstream linears. The rotated checkpoint is functionally identical (fp32 logits agree to ~3e-7, unit-gated at 1e-4) but has flatter activation/weight distributions, and remains a vanilla HF checkpoint: every existing quant config, calibrator, exporter, and runtime works on it unchanged, so the transform is orthogonal to qformat by construction. Includes the arch-mapping registry (one dict entry per supported architecture), Paley/Walsh Hadamard constructions, an external-matrix fold path gated on orthogonality (for learned rotations), bitwise seed-path reproducibility, and scoped ruff per-file-ignores for the module's math notation (mirroring the existing kernels/* exemption). Signed-off-by: Jie Ren --- .../torch/quantization/rotation/__init__.py | 18 + modelopt/torch/quantization/rotation/fold.py | 608 ++++++++++++++++++ pyproject.toml | 2 + .../torch/quantization/test_rotation_fold.py | 259 ++++++++ 4 files changed, 887 insertions(+) create mode 100644 modelopt/torch/quantization/rotation/__init__.py create mode 100644 modelopt/torch/quantization/rotation/fold.py create mode 100644 tests/unit/torch/quantization/test_rotation_fold.py diff --git a/modelopt/torch/quantization/rotation/__init__.py b/modelopt/torch/quantization/rotation/__init__.py new file mode 100644 index 00000000000..a69e552ec72 --- /dev/null +++ b/modelopt/torch/quantization/rotation/__init__.py @@ -0,0 +1,18 @@ +# 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 * diff --git a/modelopt/torch/quantization/rotation/fold.py b/modelopt/torch/quantization/rotation/fold.py new file mode 100644 index 00000000000..b14da78cfb3 --- /dev/null +++ b/modelopt/torch/quantization/rotation/fold.py @@ -0,0 +1,608 @@ +# 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). + """ + R = torch.as_tensor(mat).detach().to(torch.float64).cpu() + 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 + 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 seams of the T14 prefold, 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 + (T14 convention; default 256 — the fp16-endpoint-safe value from the T15 + 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/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_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) From a7b671c354f1873a1f6abf328c0a934a5cb9cde3 Mon Sep 17 00:00:00 2001 From: Jie Ren Date: Mon, 3 Aug 2026 15:15:41 -0700 Subject: [PATCH 2/5] Port the SGDG Cayley-SGD Stiefel-manifold optimizer (Li et al., MIT) for rotation learning Self-contained port of stiefel_optimizer.py from the MIT-licensed reference implementation of Li et al. (ICLR 2020), with the repository's helper functions inlined. Verified bitwise against the copy vendored in Meta's SpinQuant (train_utils/optimizer.py, whose own header credits Li's repository as origin) on the stiefel branch at momentum 0.0 and 0.9. Documented original quirks (inert-momentum dead store, global-random QR-retraction draw) are reproduced for trajectory parity with published SpinQuant training runs. Per the CONTRIBUTING copied-code policy: source link with commit hash and the original MIT notice precede the NVIDIA header (SPDX: Apache-2.0 AND MIT), Jun Li is added to the LICENSE third-party MIT notices, and the file is excluded from the insert-license hook. Signed-off-by: Jie Ren --- .pre-commit-config.yaml | 1 + LICENSE | 1 + modelopt/torch/quantization/rotation/sgdg.py | 219 +++++++++++++++++++ 3 files changed, 221 insertions(+) create mode 100644 modelopt/torch/quantization/rotation/sgdg.py 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/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/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 From 501b9e5f86e009f88590fb8710e43a3cb8f5b088 Mon Sep 17 00:00:00 2001 From: Jie Ren Date: Mon, 3 Aug 2026 15:15:42 -0700 Subject: [PATCH 3/5] Add learn_rotations: Cayley-SGD learned R1/R2 with pluggable fake-quant objectives learn_rotations minimizes next-token CE of the fake-quantized rotated model (SpinQuant, arXiv:2405.16406) over R1/R2 on the Stiefel manifold, reusing fold.py's architecture registry and orientation table. Each step assembles the rotated effective weights out of place and reparametrizes a frozen model, so weights are never rewritten; gradients reach only the rotation parameters (asserted at step 0). QuantObjective presets cover W4A4 per-group-128, INT8 per-tensor-static, and the paper-protocol W16A4 asym + in-graph-R4 objective; options include per-token asymmetric activation numerics (official-SpinQuant ActQuantizer parity), a KD loss against a frozen teacher, and OSTQuant-style learned seam diagonals (transform-QAT) baked by fold_seam_diags. A final polar retraction returns matrices orthogonal to ~1e-14; RotationSet save/load round-trips bitwise and refuses off-manifold matrices. steps=0 reproduces fold_rotations' seeded draws bitwise, so trained and random rotations share one provenance contract. Adds 54 CPU unit tests (61 total with the fold suite). Signed-off-by: Jie Ren --- .../torch/quantization/rotation/__init__.py | 1 + modelopt/torch/quantization/rotation/learn.py | 1042 +++++++++++++++++ .../quantization/test_rotation_ext_fold.py | 619 ++++++++++ .../quantization/test_rotation_ext_learner.py | 656 +++++++++++ .../quantization/test_rotation_ext_sgdg.py | 268 +++++ .../torch/quantization/test_rotation_kd.py | 145 +++ .../torch/quantization/test_rotation_learn.py | 481 ++++++++ .../test_rotation_paper_objective.py | 250 ++++ .../test_rotation_transform_qat.py | 708 +++++++++++ 9 files changed, 4170 insertions(+) create mode 100644 modelopt/torch/quantization/rotation/learn.py create mode 100644 tests/unit/torch/quantization/test_rotation_ext_fold.py create mode 100644 tests/unit/torch/quantization/test_rotation_ext_learner.py create mode 100644 tests/unit/torch/quantization/test_rotation_ext_sgdg.py create mode 100644 tests/unit/torch/quantization/test_rotation_kd.py create mode 100644 tests/unit/torch/quantization/test_rotation_learn.py create mode 100644 tests/unit/torch/quantization/test_rotation_paper_objective.py create mode 100644 tests/unit/torch/quantization/test_rotation_transform_qat.py diff --git a/modelopt/torch/quantization/rotation/__init__.py b/modelopt/torch/quantization/rotation/__init__.py index a69e552ec72..ec4e1c4c4b9 100644 --- a/modelopt/torch/quantization/rotation/__init__.py +++ b/modelopt/torch/quantization/rotation/__init__.py @@ -16,3 +16,4 @@ """Rotation folding + learning (SpinQuant/QuaRot R1 + R2) as pre-quantization transforms.""" from .fold import * +from .learn import * diff --git a/modelopt/torch/quantization/rotation/learn.py b/modelopt/torch/quantization/rotation/learn.py new file mode 100644 index 00000000000..7bd8139b88e --- /dev/null +++ b/modelopt/torch/quantization/rotation/learn.py @@ -0,0 +1,1042 @@ +# 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 T14 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``); T12's arm4 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): + if self.a_bits is not None and 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 internal reference runs / 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 T12's arm4 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") + self.rotations = { + k: torch.as_tensor(v).detach().to(torch.float64).cpu() + 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)}" + ) + norm[int(k)] = { + kk: torch.as_tensor(vv).detach().to(torch.float64).cpu().flatten() + 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 + T14-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) -> torch.Tensor: + if isinstance(batch, torch.Tensor): + ids = batch + elif isinstance(batch, Mapping) or hasattr(batch, "keys"): + ids = batch["input_ids"] + 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)}" + return ids + + +# -------------------------------------------------------------------------------------- +# 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 (reference budget: 150). ``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 (T25): 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 pre-T25 trainer. + kd_alpha: KD mixing weight (only with ``teacher``; T22's measured setting 0.5). + kd_temp: KD softmax temperature (only with ``teacher``; T22 setting 2.0). + + 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] + + 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) + err = (R64 @ R64.T - eye).abs().max().item() + assert err < LEARNED_ORTHO_TOL, ( + f"init_rotations[{k!r}]: not orthogonal (max |R R^T - I| = {err:.3e})" + ) + 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" + ) + + hooks = None + 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), + ) + n_hooked = hooks.attach(model) + expected_hooks = ( + n_layers * len(_ATTN_PROJS + _MLP_PROJS) + if objective_cfg.a_bits is not None + else n_layers + ) + assert n_hooked == expected_hooks, ( + f"activation hooks attached to {n_hooked} linears, expected {expected_hooks}" + ) + + 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: + 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 = _batch_input_ids(batch).to(device) + 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, labels=ids, use_cache=False) + loss = out.loss + if teacher is not None: # KD objective (T25); teacher never touched + with torch.no_grad(): + tlogits = teacher(input_ids=ids).logits + T = kd_temp + kd = torch.nn.functional.kl_div( + torch.nn.functional.log_softmax(out.logits / T, dim=-1), + torch.nn.functional.softmax(tlogits / 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/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..cf47c002d5a --- /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 (T20.1) 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..859315e406e --- /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 (T20.4) 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(AssertionError, 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_kd.py b/tests/unit/torch/quantization/test_rotation_kd.py new file mode 100644 index 00000000000..17a6b7ec825 --- /dev/null +++ b/tests/unit/torch/quantization/test_rotation_kd.py @@ -0,0 +1,145 @@ +# 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 (T25): 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(): + torch.manual_seed(11) + batch = _batches(n=1) + model = _tiny_llama() + teacher = copy.deepcopy(model) + rs = learn_rotations( + model, + batch, + steps=10, + lr=0.5, + objective_cfg=TINY_W4A4, + seed=5, + log_every=0, + teacher=teacher, + ) + assert rs.history[-1]["loss"] < rs.history[0]["loss"], ( + f"KD loss did not decrease: {rs.history[0]['loss']} -> {rs.history[-1]['loss']}" + ) + assert max(rs.ortho_audit().values()) < 1e-4 + + +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..63c35e28964 --- /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 (T26): 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..4ed5d2fabfb --- /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 (T22.3): 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): + """VERBATIM copy of the pre-T22.3 _assemble_effective_weights body — 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 pre-T22.3 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) From adac2f4563771b820326d3b92515d3b0f9bbc8d6 Mon Sep 17 00:00:00 2001 From: Jie Ren Date: Mon, 3 Aug 2026 15:15:43 -0700 Subject: [PATCH 4/5] Add rotation module design notes and changelog entry README.md documents the module's design: goal and non-goals (offline R1/R2 only; online R3/R4 out of scope), the reader/writer orientation table, the arch-mapping registry, equivalence gates, the external-matrix fold path, objective presets, and the design lineage vs the official SpinQuant trainer, including measured accuracy anchors (learned no-had W4A4 GPTQ on Llama-3.2-1B reproduces the paper's Table-8 result: 48.89 vs 48.4 WikiText-2 PPL on the official released harness). Signed-off-by: Jie Ren --- CHANGELOG.rst | 1 + .../torch/quantization/rotation/README.md | 257 ++++++++++++++++++ 2 files changed, 258 insertions(+) create mode 100644 modelopt/torch/quantization/rotation/README.md 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/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. From b6480549701b74008376a5a05ef6d4039a275dfa Mon Sep 17 00:00:00 2001 From: Jie Ren Date: Fri, 7 Aug 2026 15:23:16 -0700 Subject: [PATCH 5/5] Fix silent-failure contracts in the rotation module found by review Nine defects that produced wrong numbers, mutated caller state, or dropped user input WITHOUT raising - the class of bug the module's equivalence gates cannot catch. Each is pinned by a new test in test_rotation_contracts.py, verified to fail before the fix and pass after. Correctness: - KD term: kl_div(reduction='batchmean') on [bs, seq, vocab] logits divided by batch size alone, making the KD loss seq_len times the per-token KL, so the documented (1-alpha)*CE + alpha*T^2*KL mix silently depended on calibration sequence length. Flatten to [tokens, vocab] and exclude padded positions. - teacher=model is now rejected: the teacher forward runs inside the student's reparametrization with its hooks attached, so a self-teacher gave KL == 0 and a plain-CE run scaled by (1-alpha) while meta reported KD as active. - attention_mask is extracted from dict batches and honored: padded batches attended over padding and computed cross-entropy on pad labels, so the objective depended on how much padding the tokenizer added. - External rotations are copied, not aliased: every conversion in the accept path is a no-op for float64 CPU input, so the audited matrix tracked the caller's buffer and could change after the orthogonality gate passed. - Duplicate R2/seam-diag layer keys (int 0 and 'model.layers.0...') now raise instead of collapsing last-writer-wins and silently dropping one matrix. Validation and state hygiene: - Bit-widths below 2 are rejected: b=1 gives scale = amax/0 = inf and 0*inf, NaN-ing the entire objective with no diagnostic. - kd_temp <= 0 and kd_alpha outside [0, 1] are rejected at the call instead of surfacing as a NaN in the closing SVD (or silently maximizing CE). - a_mode is validated even when a_bits is None, so a typo cannot lie dormant. - The warm-start gate checks both orthogonality residual forms (the exit audit already did), and raises ValueError like fold_rotations rather than asserting - bare asserts vanish under python -O. - Activation hooks attach inside the guarded region: a setup failure previously left them on the caller's model, silently quantizing every later forward. Also replaces the KD training test's 'final loss < first loss' assertion, which passed on seed luck (plain CE on the same recipe also ends higher than it starts at this toy scale) and only held because the inflated KD term moved the loss; it now asserts the rotations move, stay orthogonal, and that the KD term is genuinely nonzero. Public docstrings drop internal experiment identifiers that no upstream reader can resolve. Signed-off-by: Jie Ren --- modelopt/torch/quantization/rotation/fold.py | 15 +- modelopt/torch/quantization/rotation/learn.py | 139 +++++-- .../quantization/test_rotation_contracts.py | 382 ++++++++++++++++++ .../quantization/test_rotation_ext_fold.py | 2 +- .../quantization/test_rotation_ext_learner.py | 4 +- .../torch/quantization/test_rotation_kd.py | 36 +- .../test_rotation_paper_objective.py | 2 +- .../test_rotation_transform_qat.py | 6 +- 8 files changed, 538 insertions(+), 48 deletions(-) create mode 100644 tests/unit/torch/quantization/test_rotation_contracts.py diff --git a/modelopt/torch/quantization/rotation/fold.py b/modelopt/torch/quantization/rotation/fold.py index b14da78cfb3..89ff032ffc0 100644 --- a/modelopt/torch/quantization/rotation/fold.py +++ b/modelopt/torch/quantization/rotation/fold.py @@ -157,9 +157,11 @@ 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). + 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() + 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() @@ -191,6 +193,11 @@ def _normalize_external_r2(R2, n_layers: int, head_dim: int) -> list[torch.Tenso 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( @@ -496,7 +503,7 @@ def fold_seam_diags(model: nn.Module, seam_diags, smax: float = 256.0) -> dict: 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 seams of the T14 prefold, as exact per-seam + ROTATION-SURVIVING SmoothQuant-style prefold seams, as exact per-seam functional identities — - **down seam** (``s_down [intermediate_size]``): ``up_proj`` rows ``/= s_down`` @@ -522,7 +529,7 @@ def fold_seam_diags(model: nn.Module, seam_diags, smax: float = 256.0) -> dict: 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 - (T14 convention; default 256 — the fp16-endpoint-safe value from the T15 + (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. diff --git a/modelopt/torch/quantization/rotation/learn.py b/modelopt/torch/quantization/rotation/learn.py index 7bd8139b88e..8449d289630 100644 --- a/modelopt/torch/quantization/rotation/learn.py +++ b/modelopt/torch/quantization/rotation/learn.py @@ -25,8 +25,8 @@ 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 +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``. @@ -124,7 +124,7 @@ class QuantObjective: 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 T14 SmoothQuant + 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 @@ -143,7 +143,7 @@ class QuantObjective: (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``); T12's arm4 measured that training WITHOUT it + (``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. @@ -160,10 +160,17 @@ class QuantObjective: r4_in_graph: bool = False def __post_init__(self): - if self.a_bits is not None and self.a_mode not in ( - "per_token_dynamic", - "per_tensor_static", - ): + # 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}") @@ -174,7 +181,7 @@ def __post_init__(self): #: SpinQuant-paper-style W4A4 (per-group-128 sym weights + per-token dynamic sym int4 -#: activations). Comparison point for the internal reference runs / the paper's W4A4 rows. +#: 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" ) @@ -191,7 +198,7 @@ def __post_init__(self): #: 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 T12's arm4 measured that +#: 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( @@ -398,8 +405,11 @@ class RotationSet: 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() + 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: @@ -409,8 +419,10 @@ def __post_init__(self): 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() + 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(): @@ -518,7 +530,7 @@ def _assemble_effective_weights( ``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 - T14-prefold structure BEFORE weight fake-quant, prefold-inside / rotation-outside + 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 @@ -611,11 +623,19 @@ def _iter_batches(calib_loader: Iterable, steps: int): n += 1 -def _batch_input_ids(batch) -> torch.Tensor: +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 " @@ -624,7 +644,14 @@ def _batch_input_ids(batch) -> torch.Tensor: if ids.dim() == 1: ids = ids.unsqueeze(0) assert ids.dim() == 2, f"input_ids must be [bs, seq], got shape {tuple(ids.shape)}" - return ids + 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 # -------------------------------------------------------------------------------------- @@ -671,7 +698,8 @@ def learn_rotations( 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 (reference budget: 150). ``steps=0`` returns + 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"`` @@ -684,14 +712,14 @@ def learn_rotations( 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 (T25): loss becomes + 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 pre-T25 trainer. - kd_alpha: KD mixing weight (only with ``teacher``; T22's measured setting 0.5). - kd_temp: KD softmax temperature (only with ``teacher``; T22 setting 2.0). + 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 @@ -713,6 +741,21 @@ def learn_rotations( ) 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 @@ -776,10 +819,19 @@ def learn_rotations( ) for k, R64 in draws.items(): eye = torch.eye(R64.shape[0], dtype=torch.float64) - err = (R64 @ R64.T - eye).abs().max().item() - assert err < LEARNED_ORTHO_TOL, ( - f"init_rotations[{k!r}]: not orthogonal (max |R R^T - I| = {err:.3e})" + # 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)) @@ -836,21 +888,21 @@ def learn_rotations( "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), ) - n_hooked = hooks.attach(model) expected_hooks = ( n_layers * len(_ATTN_PROJS + _MLP_PROJS) if objective_cfg.a_bits is not None else n_layers ) - assert n_hooked == expected_hooks, ( - f"activation hooks attached to {n_hooked} linears, expected {expected_hooks}" - ) if learn_seam_diag: assert base["model.layers.0.mlp.up_proj.weight"].shape[0] == intermediate, ( @@ -883,6 +935,11 @@ def _r1_ortho() -> float: # 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 @@ -891,7 +948,14 @@ def _r1_ortho() -> float: 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 = _batch_input_ids(batch).to(device) + 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, @@ -908,15 +972,24 @@ def _r1_ortho() -> float: # 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, labels=ids, use_cache=False) + out = model(input_ids=ids, attention_mask=attn_mask, labels=labels, use_cache=False) loss = out.loss - if teacher is not None: # KD objective (T25); teacher never touched + if teacher is not None: # KD objective; teacher never touched with torch.no_grad(): - tlogits = teacher(input_ids=ids).logits + 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(out.logits / T, dim=-1), - torch.nn.functional.softmax(tlogits / T, dim=-1), + 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 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 index cf47c002d5a..a687b86edc4 100644 --- a/tests/unit/torch/quantization/test_rotation_ext_fold.py +++ b/tests/unit/torch/quantization/test_rotation_ext_fold.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Extensive fold property tests (T20.1) for +"""Extensive fold property tests for modelopt.torch.quantization.rotation.fold_rotations. Coverage beyond test_rotation_fold.py: diff --git a/tests/unit/torch/quantization/test_rotation_ext_learner.py b/tests/unit/torch/quantization/test_rotation_ext_learner.py index 859315e406e..5b9aadbebd1 100644 --- a/tests/unit/torch/quantization/test_rotation_ext_learner.py +++ b/tests/unit/torch/quantization/test_rotation_ext_learner.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Learner-semantics tests (T20.4) for modelopt.torch.quantization.rotation.learn. +"""Learner-semantics tests for modelopt.torch.quantization.rotation.learn. Covers: 1. steps=0 == fold_rotations seed draws, bitwise, seeds {0, 3, 7}, both archs. @@ -488,7 +488,7 @@ def test_warm_start_continues_donor_loss_and_gate_rejects_corruption(): # Ortho gate: corrupted warm start must be refused. bad = dict(donor.rotations) bad["R1"] = bad["R1"] * 1.01 - with pytest.raises(AssertionError, match="not orthogonal"): + with pytest.raises(ValueError, match="not orthogonal"): learn_rotations( _tiny_qwen3(), batch, diff --git a/tests/unit/torch/quantization/test_rotation_kd.py b/tests/unit/torch/quantization/test_rotation_kd.py index 17a6b7ec825..96cf758d646 100644 --- a/tests/unit/torch/quantization/test_rotation_kd.py +++ b/tests/unit/torch/quantization/test_rotation_kd.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for the KD rotation objective (T25): learn_rotations(teacher=...).""" +"""Tests for the KD rotation objective: learn_rotations(teacher=...).""" import copy import sys @@ -93,10 +93,24 @@ def test_kd_wiring_no_quant_kl_vanishes(): 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, @@ -107,10 +121,24 @@ def test_kd_objective_trains_and_stays_orthogonal(): log_every=0, teacher=teacher, ) - assert rs.history[-1]["loss"] < rs.history[0]["loss"], ( - f"KD loss did not decrease: {rs.history[0]['loss']} -> {rs.history[-1]['loss']}" - ) + # (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(): diff --git a/tests/unit/torch/quantization/test_rotation_paper_objective.py b/tests/unit/torch/quantization/test_rotation_paper_objective.py index 63c35e28964..04d578c69b8 100644 --- a/tests/unit/torch/quantization/test_rotation_paper_objective.py +++ b/tests/unit/torch/quantization/test_rotation_paper_objective.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for the paper-protocol objective extensions (T26): per-token ASYM min-max +"""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``). diff --git a/tests/unit/torch/quantization/test_rotation_transform_qat.py b/tests/unit/torch/quantization/test_rotation_transform_qat.py index 4ed5d2fabfb..613bbd64790 100644 --- a/tests/unit/torch/quantization/test_rotation_transform_qat.py +++ b/tests/unit/torch/quantization/test_rotation_transform_qat.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Transform-QAT tests (T22.3): jointly learned rotations R1/R2 + per-input-channel seam +"""Transform-QAT tests: jointly learned rotations R1/R2 + per-input-channel seam diagonals (OSTQuant-style) in modelopt.torch.quantization.rotation. Covers: @@ -186,7 +186,7 @@ def _assembled_logits(model, eff): def _assemble_pre_change(base, R1, R2s, n_layers, head_dim, objective, out_dtype): - """VERBATIM copy of the pre-T22.3 _assemble_effective_weights body — the bitwise + """Rotation-only reference assembly (no seam diagonals) — the bitwise oracle for the learn_seam_diag=False path.""" compute = R1.dtype d = head_dim @@ -627,7 +627,7 @@ def test_save_load_roundtrip_with_seam_diags_and_old_format(): 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 pre-T22.3 R.bin): loads fine, + # 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)