Add asp_nvfp4_runtime.py (NVFP4 deflated ASP)

asp_nvfp4_runtime.py

@@ -0,0 +1,195 @@
+"""Run the deflated-ASP NVFP4 checkpoint. Real 4-bit execution on the FP4 tensor cores.
+
+Per layer, exactly the contract the export wrote and the search scored:
+
+    x~ = (x / s) H
+    y  = (x~ V)(W~V)^T  +  NVFP4GEMM( (I - VV^T) x~ ,  W~(I-VV^T) )   + bias
+
+The weight never leaves E2M1: it is packed nibbles with pre-swizzled E4M3 block scales, handed
+straight to `torch._scaled_mm_v2` with recipe BlockWise1x16. Nothing is dequantised to bf16 and no
+fp16 GEMM runs on the bulk path -- that is what makes this real quantisation rather than a
+simulation of it.
+
+The ACTIVATION is quantised per forward: block 16, E4M3 scales, one level (SVDQuant's nvfp4 recipe
+for inputs), then swizzled. The swizzle is unavoidable -- `_scaled_mm_v2` rejects row-major scales
+-- and it is done in torch here rather than fused, so this path is correct first and fast later.
+The rank-r branch stays in bf16 by design: it carries the directions the action depends on, which
+is the whole point of protecting them.
+
+  FW_ASP_CKPT=/path/to/ur3_step7000_asp_nvfp4.pt
+"""
+from __future__ import annotations
+
+import math
+import os
+
+import torch
+import torch.nn as nn
+
+from nvfp4 import nvfp4_mm, quantize_nvfp4, recip_scale, swizzle_scales
+
+
+def _fwht(x):
+    """Orthonormal fast Walsh-Hadamard transform along the last dim (size a power of 2)."""
+    orig = x.shape; m = orig[-1]; x = x.reshape(-1, m).clone(); h = 1
+    while h < m:
+        x = x.view(-1, m // (2 * h), 2, h); a = x[:, :, 0, :]; b = x[:, :, 1, :]
+        x = torch.stack([a + b, a - b], dim=2).reshape(-1, m); h *= 2
+    return (x / math.sqrt(m)).reshape(orig)
+
+
+_HMAT_CACHE = {}
+
+
+def _hadamard_matrix(B, dtype, device):
+    """H with x @ H == _fwht(x). Built by running _fwht on the identity, so the butterfly
+    ORDERING and the 1/sqrt(B) normalisation match the ones the weights were rotated with at
+    export. A Sylvester construction would give a different ordering, and the resulting error
+    would look like a quantisation regression rather than a transform mismatch."""
+    key = (B, dtype, str(device))
+    if key not in _HMAT_CACHE:
+        eye = torch.eye(B, device=device, dtype=torch.float32)
+        _HMAT_CACHE[key] = _fwht(eye).to(dtype).contiguous()
+    return _HMAT_CACHE[key]
+
+
+_H: dict = {}
+
+
+def _hmat(B, device, dtype):
+    k = (B, str(device), dtype)
+    if k not in _H:
+        _H[k] = _hadamard_matrix(B, dtype, device)
+    return _H[k]
+
+
+class ASPNVFP4Linear(nn.Module):
+    """One exported Linear: NVFP4 deflated weight + bf16 rank-r action subspace."""
+
+    def __init__(self, e: dict, dtype, device):
+        super().__init__()
+        self.in_features = int(e["in_features"])
+        self.out_features = int(e["out_features"])
+        self.block = int(e["fwht_block"])
+        self.a_block = int(e["a_block"])
+        self.rank = int(e["asp_rank"])
+        self.register_buffer("wq", e["wq"].to(device))
+
+        self.register_buffer("wsc", e["wscale_swizzled"].to(device))
+        # The kernel wants the RECIPROCAL of the per-tensor scale, and it is a constant: built
+        # once here, not per forward. The transposed weight view is hoisted for the same reason.
+        self.register_buffer("w_rglob", recip_scale(e["wglobal"].to(device), device))
+        self.register_buffer("a_rglob", recip_scale(None, device).clone())
+        # KEPT IN FP32, NOT THE MODEL DTYPE. These are stored fp16; casting them to bf16 throws
+        # away three mantissa bits before a 4-bit grid ever sees them, and the grid's codes are
+        # ~33% apart, so those bits decide which code an element lands on. Measured: bf16 buffers
+        # put the layer output 2e-2 to 5e-2 from the fp32 reference; fp16 -> fp32 is exact and
+        # costs a few hundred MB across the model for tensors that are r=32 columns wide.
+        self.register_buffer("inv_smooth", e["inv_smooth"].to(device=device, dtype=torch.float32))
+        self.register_buffer("bias", e["bias"].to(device=device, dtype=torch.float32)
+                             if "bias" in e else None)
+        # V as [in, r] for the projection, W~V as [out, r] for the epilogue
+        self.register_buffer("V", e["lr_a"].t().contiguous().to(device=device,
+                                                                dtype=torch.float32)
+                             if self.rank else None)
+        self.register_buffer("WV", e["lr_b"].to(device=device, dtype=torch.float32)
+                             if self.rank else None)
+        self._wqt = self.wq.t()          # transposed view, hoisted out of the forward
+
+    def forward(self, x):
+        # THE PROLOGUE RUNS IN FP32, deliberately. What follows it is a 4-bit grid whose codes are
+        # ~33% apart, so a bf16 rounding of the smoothed/rotated activation moves elements across
+        # code boundaries: measured on the exported layers, a bf16 prologue put the layer output
+        # 3.8e-2 to 9.0e-2 from the fp32 one. The search that CHOSE this layer's (alpha, beta) and
+        # its subspace scored an fp32 prologue, so a bf16 one would deploy a different arm than the
+        # one that was selected. It is also cheap: the rotation is O(d*sqrt(d)) beside an FP4 GEMM.
+        shp = x.shape
+        x2 = x.reshape(-1, shp[-1]).float()
+        xh = x2 * self.inv_smooth
+        D = xh.shape[-1]
+        xt = (xh.reshape(-1, D // self.block, self.block)
+              @ _hmat(self.block, xh.device, torch.float32)).reshape(-1, D)
+        if self.rank:
+            Vf = self.V
+            proj = xt @ Vf                              # [tok, r]
+            xp = xt - proj @ Vf.t()                     # (I - VV^T) x~
+        else:
+            proj, xp = None, xt
+        ak, asc, _ = quantize_nvfp4(xp, block=self.a_block, two_level=False)
+        y = nvfp4_mm(ak, swizzle_scales(asc), self.wq, self.wsc, self.a_rglob, self.w_rglob,
+                     out_dtype=torch.float32, b_packed_t=self._wqt)[: xp.shape[0]]
+        # EVERYTHING IN FP32, THEN ONE CAST AT THE END. The epilogue and the bias are fp32
+        # buffers, so casting before adding them promotes the result straight back to fp32 and the
+        # next layer_norm fails on a dtype it did not expect. One cast, last.
+        if self.rank:
+            y = y + proj @ self.WV.t()
+        if self.bias is not None:
+            y = y + self.bias
+        y = y.to(x.dtype)
+        return y.reshape(*shp[:-1], self.out_features)
+
+    def extra_repr(self):
+        return (f"in={self.in_features}, out={self.out_features}, NVFP4 w{self.a_block}, "
+                f"rank={self.rank}, H={self.block}")
+
+
+def install_asp_nvfp4(model, ckpt_path=None, verbose=True) -> int:
+    """Swap every exported Linear. A partial swap is not a defined arm, so it raises."""
+    ckpt_path = ckpt_path or os.environ.get("FW_ASP_CKPT")
+    blob = torch.load(str(ckpt_path), map_location="cpu", weights_only=False)
+    meta, layers = blob["meta"], blob["layers"]
+    if meta.get("lowrank_mode") != "asp_deflated":
+        raise RuntimeError(f"{ckpt_path}: lowrank_mode={meta.get('lowrank_mode')!r}; this runtime "
+                           f"implements the DEFLATED ASP contract and applying it to another "
+                           f"would give a well-formed GEMM of the wrong bilinear form.")
+    if verbose:
+        f = meta["format"]
+        print(f"[asp-nvfp4] {meta['scheme']}: {meta['n_layers']} Linears "
+              f"({meta['asp_layers']} with ASP r{meta['rank']}), {meta['bpw']} BPW", flush=True)
+        print(f"[asp-nvfp4]   {f['element']} w{f['w_block']}/a{f['a_block']}, {f['scale']} scales, "
+              f"{f['weight_scale_levels']}-level weights, {f['scale_layout']}", flush=True)
+    named = dict(model.named_modules())
+    done, missing = 0, []
+    for name, e in layers.items():
+        mod = named.get(name)
+        if not isinstance(mod, nn.Linear):
+            missing.append(name)
+            continue
+        new = ASPNVFP4Linear(e, mod.weight.dtype, mod.weight.device)
+        parent = model.get_submodule(name.rsplit(".", 1)[0])
+        setattr(parent, name.rsplit(".", 1)[-1], new)
+        named[name] = None
+        del mod, new
+        done += 1
+    if missing:
+        raise RuntimeError(f"{ckpt_path}: {len(missing)} layers did not resolve, e.g. {missing[:3]}")
+    del blob, layers, named
+    import gc
+    gc.collect()
+    torch.cuda.empty_cache()
+    if verbose:
+        free, tot = torch.cuda.mem_get_info()
+        print(f"[asp-nvfp4] replaced {done} Linears | GPU {(tot-free)/2**30:.1f}/"
+              f"{tot/2**30:.0f} GiB", flush=True)
+    return done
+
+
+def load_quantized_asp_model(ckpt_path, build_model, verbose=True):
+    """Build the model straight from the self-contained NVFP4 checkpoint.
+
+    `build_model` is your own bf16 FastWAM constructor and must return `(model, cfg)`; the
+    checkpoint carries every tensor the quantised model needs, so it is only used for the
+    module graph and the non-Linear submodules. Nothing is read from the bf16 weights.
+    """
+    blob = torch.load(str(ckpt_path), map_location="cpu", weights_only=False)
+    if "mot_rest" not in blob:
+        raise RuntimeError(f"{ckpt_path} is not self-contained")
+    model, cfg = build_model()
+    missing, unexpected = model.mot.load_state_dict(blob["mot_rest"], strict=False)
+    if unexpected:
+        raise RuntimeError(f"{len(unexpected)} unexpected mot tensors, e.g. {list(unexpected)[:3]}")
+    if model.proprio_encoder is not None:
+        model.proprio_encoder.load_state_dict(blob["proprio_encoder"], strict=True)
+    del blob
+    install_asp_nvfp4(model, ckpt_path, verbose=verbose)
+    return model, cfg