From f8c0552c6da5e45d1f1fa5d48e15ca04717d3602 Mon Sep 17 00:00:00 2001
From: ZacharyZcR <zacharyzcr1984@gmail.com>
Date: Tue, 14 Jul 2026 02:45:34 +0800
Subject: [PATCH] quant_ablation: add rotation preconditioning (-rot,
 QuaRot/QuIP#) and int3 schemes (#132, #81)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Extends the ablation grammar to int{2,3,4,8}[-g<N>][-rot][-nohead]. -rot round-trips weights through an orthogonal Hadamard Q=diag(±1)·H/√n on the input dim, measuring the exact weight error of a deployed rotate-activations scheme. Engine-free tool (c/tools/quant_ablation.py only). Verified: syntax clean, scheme parser correct (int3/-rot/-nohead), no unsafe constructs. Findings feed the v2 int3-g64 direction (#81/#108).
---
 c/tools/quant_ablation.py | 51 +++++++++++++++++++++++++++++++++------
 1 file changed, 44 insertions(+), 7 deletions(-)

diff --git a/c/tools/quant_ablation.py b/c/tools/quant_ablation.py
index c3c5bf9..dd192ba 100644
--- a/c/tools/quant_ablation.py
+++ b/c/tools/quant_ablation.py
@@ -78,14 +78,51 @@ def _quant_last_dim(x, bits, group):
     return out.reshape(*out.shape[:-2], -1) if group else out
 
 
-def quantize_param(w, bits, group):
+# --------------------------------------------------------------------------------------
+# Rotation preconditioning (QuaRot / QuIP# family, #81): multiply the input dimension by
+# an orthogonal Q = diag(signs) @ H/sqrt(n) BEFORE quantizing, and by Q^T after — the
+# round-trip Q4(W@Q)@Q.T measures exactly the weight error of a deployed scheme that
+# stores W@Q quantized and rotates activations at runtime (W'@x' = W@x since Q@Q.T = I;
+# the runtime cost is one O(D log D) transform per matmul INPUT, not per weight).
+# Spreading outliers across the block is the point: absmax scales stop being hostage to
+# one heavy coordinate, which is the failure mode #108 measured (margin erosion on MMLU).
+# --------------------------------------------------------------------------------------
+_ROT_CACHE = {}
+
+def rotation(dim, device, seed=417):
+    key = (dim, str(device))
+    if key in _ROT_CACHE:
+        return _ROT_CACHE[key]
+    if dim & (dim - 1):
+        raise SystemExit(f"-rot needs power-of-2 input dims (got {dim}); OLMoE dims are 2048/1024")
+    h = torch.ones(1, 1, device=device, dtype=torch.float32)
+    while h.shape[0] < dim:                       # Sylvester recursion
+        h = torch.cat([torch.cat([h, h], 1), torch.cat([h, -h], 1)], 0)
+    h /= h.shape[0] ** 0.5                        # orthonormal
+    g = torch.Generator().manual_seed(seed + dim)
+    signs = (torch.randint(0, 2, (dim,), generator=g).float() * 2 - 1).to(device)
+    q = signs[:, None] * h                        # Q = D @ H/sqrt(n), orthogonal
+    _ROT_CACHE[key] = q
+    return q
+
+
+def quantize_param(w, bits, group, rot=False):
     if w.ndim == 3:                        # fused experts [E, in, out] -> move input last
         x = w.transpose(1, 2).contiguous()
-        return _quant_last_dim(x, bits, group).transpose(1, 2).contiguous()
+        x = _rot_quant(x, bits, group) if rot else _quant_last_dim(x, bits, group)
+        return x.transpose(1, 2).contiguous()
+    if rot:
+        return _rot_quant(w, bits, group)
     return _quant_last_dim(w, bits, group)  # nn.Linear [out, in] -- input already last
 
 
-SCHEME_RE = re.compile(r"^int(2|4|8)(?:-g(\d+))?(-nohead)?$")
+def _rot_quant(x, bits, group):
+    """W -> Qn(W@Q) @ Q^T along the last (input) dim — see rotation() above."""
+    q = rotation(x.shape[-1], x.device)
+    return (_quant_last_dim(x.float() @ q, bits, group) @ q.T).contiguous()
+
+
+SCHEME_RE = re.compile(r"^int(2|3|4|8)(?:-g(\d+))?(-rot)?(-nohead)?$")
 
 
 def parse_scheme(name):
@@ -94,8 +131,8 @@ def parse_scheme(name):
         return None
     m = SCHEME_RE.match(name)
     if not m:
-        raise SystemExit(f"bad scheme '{name}' (expected fp16 | int{{2,4,8}}[-g<N>][-nohead])")
-    return int(m.group(1)), int(m.group(2) or 0), bool(m.group(3))
+        raise SystemExit(f"bad scheme '{name}' (expected fp16 | int{{2,3,4,8}}[-g<N>][-rot][-nohead])")
+    return int(m.group(1)), int(m.group(2) or 0), bool(m.group(3)), bool(m.group(4))
 
 
 def is_router(name):
@@ -115,7 +152,7 @@ def apply_scheme(model, scheme):
     spec = parse_scheme(scheme)
     if spec is None:
         return 0, 0, total
-    bits, group, skip_head = spec
+    bits, group, rot, skip_head = spec
     n = qp = 0
     with torch.no_grad():
         for name, p in model.named_parameters():
@@ -123,7 +160,7 @@ def apply_scheme(model, scheme):
                 continue
             if skip_head and is_head_or_embed(name):
                 continue
-            p.data.copy_(quantize_param(p.data.float(), bits, group).to(p.dtype))
+            p.data.copy_(quantize_param(p.data.float(), bits, group, rot).to(p.dtype))
             n += 1
             qp += p.numel()
     return n, qp, total