flux2distill/svdquant.py

"""SVDQuant (our own fake-quant implementation) for FLUX.2 klein 4B.

W4A8 post-training quantization with a high-precision low-rank branch that absorbs
outliers, following SVDQuant (Li et al., ICLR 2025; arXiv 2411.05007). This is a
*simulation* (fake-quant: quantize -> dequantize in the model dtype) used to map the
quality frontier on our A100 dev rig; real low-bit kernels / Nunchaku deployment are a
later concern. The decomposition is kernel-agnostic, so a checkpoint built here is the
same math a fused W4A8 kernel would run.

Per Linear  y = x Wᵀ + b  (PyTorch weight W: (out, in)):

  1. Smooth   migrate per-channel activation outliers into the weights:
              x̂ = x / s ,  Ŵ = W ⊙ s   (s ∈ ℝ_in, x̂ Ŵᵀ = x Wᵀ exactly).
              s_j = max|x_j|^α / max|W_:,j|^(1-α)   (SmoothQuant/AWQ form; α tunable).
              (max|x_j| from calibration.)
  2. Absorb   SVD of the smoothed weight, keep the top-`rank` singular components as a
              16-bit low-rank branch  L = (U_r Σ_r)(V_rᵀ)  that soaks up the dominant
              (outlier) energy.  Residual  R = Ŵ - L  is now smooth -> easy to 4-bit.
  3. Quantize residual weights to 4-bit (group-wise along in-dim), activations to 8-bit
              (per-token dynamic).  The low-rank branch sees full-precision x̂.

Forward:  y = (x̂ · L)            [16-bit, rank ≪ in]
            + Q8(x̂) · Q4(R)ᵀ     [low-bit residual]
            + b

Storage here is fake-quant (dequantized values kept in bf16), so this does NOT shrink the
checkpoint — it measures *quality*. `compressed_bytes()` reports the size a real low-bit
packing would reach.
"""
from __future__ import annotations

import torch
import torch.nn as nn

# Prefer torchao's affine-quant primitives (the requested stack); fall back to a
# numerically-identical pure-torch path if torchao is absent or its API has drifted.
try:
    import torchao  # noqa: F401
    _HAS_TORCHAO = True
except Exception:
    _HAS_TORCHAO = False


# --------------------------------------------------------------------------------------
# fake-quant primitives (symmetric, dequantized back to the input dtype)
# --------------------------------------------------------------------------------------
def fake_quant_act(x: torch.Tensor, bits: int = 8, group: int = 0) -> torch.Tensor:
    """Symmetric dynamic quant-dequant of activations.

    group=0: per-token (one scale per last-dim vector) — fine at 8-bit, but at 4-bit a single
    outlier channel forces the whole token's scale and flattens everything else.
    group>0: per-group along the channel dim (one dynamic scale per `group` channels per token)
    — the SVDQuant paper's actual W4A4 granularity (group 64, same as its weights)."""
    qmax = (1 << (bits - 1)) - 1                       # 127 for 8-bit, 7 for 4-bit
    if group and x.shape[-1] % group == 0:
        shp = x.shape
        xg = x.reshape(*shp[:-1], shp[-1] // group, group)
        scale = xg.detach().abs().amax(dim=-1, keepdim=True) / qmax
        scale = scale.clamp(min=1e-8)
        xq = torch.clamp(torch.round(xg / scale), -qmax - 1, qmax)
        return (xq * scale).reshape(shp).to(x.dtype)
    scale = x.detach().abs().amax(dim=-1, keepdim=True) / qmax
    scale = scale.clamp(min=1e-8)
    xq = torch.clamp(torch.round(x / scale), -qmax - 1, qmax)
    return (xq * scale).to(x.dtype)


def fake_quant_weight(W: torch.Tensor, bits: int = 4, group: int = 64) -> torch.Tensor:
    """Group-wise (along in-dim) symmetric quant-dequant of a weight (out, in)."""
    out_f, in_f = W.shape
    qmax = (1 << (bits - 1)) - 1                       # 7 for 4-bit
    g = group if (group and in_f % group == 0) else in_f
    Wg = W.reshape(out_f, in_f // g, g)
    scale = Wg.detach().abs().amax(dim=-1, keepdim=True) / qmax
    scale = scale.clamp(min=1e-8)
    Wq = torch.clamp(torch.round(Wg / scale), -qmax - 1, qmax)
    return (Wq * scale).reshape(out_f, in_f).to(W.dtype)


# --------------------------------------------------------------------------------------
# NVFP4 (Blackwell-native FP4): E2M1 elements, group-16 blocks, FP8(E4M3) per-block scales.
# This is the format the 5th-gen tensor cores accelerate (and what Nunchaku's svdq-fp4 file
# runs). Unlike symmetric-int 4-bit (15 uniform levels) the element grid is the 8 E2M1
# magnitudes {0,.5,1,1.5,2,3,4,6}·sign — non-uniform, denser near zero, but coarser in the
# bulk; the win comes from the much finer group (16 vs 64) and an FP8 (not bf16) scale.
# --------------------------------------------------------------------------------------
_E2M1_LEVELS = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
_E2M1_THRESH = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0)   # round-to-nearest midpoints
_E2M1_MAX = 6.0
_E4M3_MAX = 448.0


def _round_e2m1(a: torch.Tensor) -> torch.Tensor:
    """Round non-negative magnitudes to the nearest representable E2M1 level."""
    thr = torch.tensor(_E2M1_THRESH, device=a.device, dtype=a.dtype)
    lvl = torch.tensor(_E2M1_LEVELS, device=a.device, dtype=a.dtype)
    return lvl[torch.bucketize(a, thr)]


def fake_quant_nvfp4(x: torch.Tensor, group: int = 16) -> torch.Tensor:
    """NVFP4 fake-quant: per-group E2M1 with an FP8-E4M3 block scale + a per-tensor FP32
    global scale (NVIDIA's two-level NVFP4). Works for weights (groups along in-dim) and
    activations (groups along the channel dim). group=16 is NVFP4's native block."""
    shp = x.shape
    C = shp[-1]
    g = group if (group and C % group == 0) else C
    xg = x.reshape(*shp[:-1], C // g, g)
    t_amax = x.detach().abs().amax().clamp(min=1e-8)
    gscale = (t_amax / (_E2M1_MAX * _E4M3_MAX)).clamp(min=1e-12)        # fp32 per-tensor
    b_amax = xg.detach().abs().amax(dim=-1, keepdim=True).clamp(min=1e-8)
    bscale = b_amax / _E2M1_MAX / gscale                               # block scale, gscale units
    bscale_q = bscale.to(torch.float8_e4m3fn).to(torch.float32).clamp(min=1e-12)  # E4M3 scale
    eff = bscale_q * gscale                                            # real per-block scale
    xn = xg / eff
    xq = torch.sign(xn) * _round_e2m1(xn.abs().clamp(max=_E2M1_MAX))
    return (xq * eff).reshape(shp).to(x.dtype)


def fake_quant_fp8(x: torch.Tensor, group: int = 0) -> torch.Tensor:
    """FP8 (E4M3) dynamic fake-quant of activations — the Blackwell-native 8-bit format
    (paired with NVFP4 weights for a 'W4A8' variant that keeps activations easy)."""
    shp = x.shape
    if group and shp[-1] % group == 0:
        xg = x.reshape(*shp[:-1], shp[-1] // group, group)
        scale = (xg.detach().abs().amax(dim=-1, keepdim=True) / _E4M3_MAX).clamp(min=1e-8)
        xq = (xg / scale).to(torch.float8_e4m3fn).to(x.dtype) * scale.to(x.dtype)
        return xq.reshape(shp).to(x.dtype)
    scale = (x.detach().abs().amax(dim=-1, keepdim=True) / _E4M3_MAX).clamp(min=1e-8)
    return ((x / scale).to(torch.float8_e4m3fn).to(x.dtype) * scale.to(x.dtype)).to(x.dtype)


def _quant_weight(R: torch.Tensor, w_fmt: str, w_bits: int, w_group: int) -> torch.Tensor:
    """Dispatch the residual-weight fake-quant by format."""
    if w_fmt == "nvfp4":
        return fake_quant_nvfp4(R, group=(w_group or 16))
    return fake_quant_weight(R, w_bits, w_group)


def _quant_act(x: torch.Tensor, a_fmt: str, a_bits: int, a_group: int) -> torch.Tensor:
    """Dispatch the activation fake-quant by format."""
    if a_fmt == "nvfp4":
        return fake_quant_nvfp4(x, group=(a_group or 16))
    if a_fmt == "fp8":
        return fake_quant_fp8(x, group=a_group)
    return fake_quant_act(x, a_bits, group=a_group)


# --------------------------------------------------------------------------------------
# module
# --------------------------------------------------------------------------------------
class SVDQuantLinear(nn.Module):
    """Drop-in replacement for nn.Linear: low-rank (high-precision) + low-bit residual."""

    def __init__(self, in_f, out_f, bias, rank, w_bits=4, a_bits=8, w_group=64,
                 a_group=0, dtype=torch.bfloat16, w_fmt="int", a_fmt="int"):
        super().__init__()
        self.in_features, self.out_features = in_f, out_f
        self.rank, self.w_bits, self.a_bits, self.w_group = rank, w_bits, a_bits, w_group
        self.a_group = a_group                          # 0 = per-token act scales
        self.w_fmt, self.a_fmt = w_fmt, a_fmt           # "int" | "nvfp4" (w), + "fp8" (a)
        self.register_buffer("smooth", torch.ones(in_f, dtype=dtype))
        self.register_buffer("lora_down", torch.zeros(rank, in_f, dtype=dtype))
        self.register_buffer("lora_up", torch.zeros(out_f, rank, dtype=dtype))
        self.register_buffer("w_res", torch.zeros(out_f, in_f, dtype=dtype))  # fake-quant'd
        if bias:
            self.register_buffer("bias", torch.zeros(out_f, dtype=dtype))
        else:
            self.bias = None

    def forward(self, x):
        dt = x.dtype
        xs = x / self.smooth.to(dt)
        # high-precision low-rank branch (absorbs the outliers)
        y = (xs @ self.lora_down.to(dt).t()) @ self.lora_up.to(dt).t()
        # low-bit residual branch: quantized activations × low-bit (already fake-quant'd) weights
        xq = _quant_act(xs, self.a_fmt, self.a_bits, self.a_group)
        y = y + xq @ self.w_res.to(dt).t()
        if self.bias is not None:
            y = y + self.bias.to(dt)
        return y

    def compressed_bytes(self):
        """Size a *real* low-bit packing would reach (4-bit residual + per-group scales +
        bf16 low-rank + bf16 bias). The fake-quant buffers themselves are bf16."""
        out_f, in_f, r = self.out_features, self.in_features, self.rank
        g = self.w_group if (self.w_group and in_f % self.w_group == 0) else in_f
        res = out_f * in_f * self.w_bits / 8            # 4-bit residual
        # NVFP4 scales are FP8 (1 byte) per group-16 (+ a per-tensor fp32); int scales are bf16
        res_scale = out_f * (in_f // g) * (1 if self.w_fmt == "nvfp4" else 2)
        lowrank = (r * in_f + out_f * r) * 2            # bf16 low-rank
        smooth = in_f * 2
        bias = out_f * 2 if self.bias is not None else 0
        return res + res_scale + lowrank + smooth + bias

    @classmethod
    @torch.no_grad()
    def from_linear(cls, lin: nn.Linear, act_absmax: torch.Tensor, rank=32, alpha=0.5,
                    w_bits=4, a_bits=8, w_group=64, a_group=0, svd_device="cuda",
                    act_gram: torch.Tensor = None, whiten=True, gram_ridge=1e-3,
                    refine_iters=3, smooth=True, w_fmt="int", a_fmt="int"):
        """Build a quantized layer from a trained nn.Linear + its calibration activation stats.

        Two SVD modes for the low-rank fit:
        - plain (whiten=False): SVD of the smoothed weight Ŵ, minimizing ‖Ŵ−L‖ (the base
          SVDQuant paper's headline derivation). Absorbs the largest WEIGHT singular values.
        - activation-aware (whiten=True, needs `act_gram`): SVD in the activation metric,
          minimizing the OUTPUT error ‖X̂(Ŵ−L)‖ (ASVD / SVD-LLM style). With Ĝ=X̂ᵀX̂=M·Mᵀ
          (eigen square-root), SVD(Ŵ·M) gives the rank-r factor capturing the most output-
          relevant energy; map back by M⁻¹. Strictly stronger — makes rank a real knob.

        ITERATIVE REFINEMENT (SVDQuant §4.2): instead of fitting L once and quantizing the
        residual R=Ŵ−L once, we iterate `refine_iters` times: re-fit the low-rank branch to
        Ŵ−Q(R) so it ABSORBS the 4-bit rounding error, re-quantize, and keep the iterate with
        the smallest (metric-weighted) reconstruction error. This is the paper's error-feedback
        loop; it costs a few extra SVDs at build time and nothing at inference. refine_iters=0
        reproduces the one-shot decomposition.

        `act_gram` is the RAW activation Gram Gᵣ=XᵀX (in×in); the smoothed gram is derived as
        Ĝ = Gᵣ⊘(s·sᵀ) since smoothing is diagonal (so one calibration pass suffices).

        Returns (module, diag).
        """
        dtype = lin.weight.dtype
        W = lin.weight.data.float()                     # (out, in)
        out_f, in_f = W.shape
        dev = W.device

        # 1. smoothing factor (migrate activation outliers -> weights). smooth=False -> s=1
        #    (no SmoothQuant: raw weights to 4-bit, raw activations to 8-bit — the RTN floor).
        w_absmax = W.abs().amax(dim=0).clamp(min=1e-8)            # (in,)
        a_absmax = act_absmax.float().to(dev).clamp(min=1e-8)     # (in,)
        if smooth:
            s = (a_absmax.pow(alpha) / w_absmax.pow(1.0 - alpha)).clamp(min=1e-8)
        else:
            s = torch.ones_like(w_absmax)
        What = W * s.view(1, -1)                                  # smoothed weight (out,in)
        Wd = What.to(svd_device)

        # whitening square-root M (M Mᵀ = Ĝ) and its inverse, or None for plain SVD.
        use_whiten = whiten and act_gram is not None
        Ghat = M = Minv = None
        if use_whiten:
            # robust symmetric-eigen square-root (Cholesky is fragile: the bf16-accumulated
            # Gram has near-zero / slightly-negative directions → not numerically PD).
            try:
                G = act_gram.float().to(svd_device)              # raw Gram XᵀX (in,in)
                sd = s.to(svd_device)
                Ghat = G / sd.view(-1, 1) / sd.view(1, -1)       # smoothed-act gram X̂ᵀX̂
                Ghat = 0.5 * (Ghat + Ghat.t())                   # symmetrize
                evals, evecs = torch.linalg.eigh(Ghat)           # ascending
                emax = evals[-1].clamp(min=1e-12)
                ev = evals.clamp(min=gram_ridge * emax)          # floor tiny/neg eigs (ridge)
                sqrt_ev = ev.sqrt()
                M = (evecs * sqrt_ev) @ evecs.t()                # symmetric √Ĝ
                Minv = (evecs * (1.0 / sqrt_ev)) @ evecs.t()     # symmetric √Ĝ⁻¹
            except Exception as e:                               # any failure -> plain SVD
                import warnings
                warnings.warn(f"whitened SVD failed ({type(e).__name__}: {e}); plain-SVD fallback")
                use_whiten = False
                Ghat = M = Minv = None

        def lowrank_fit(T):
            """Rank-r factor (up,down) of target T, in the whitened metric if M is set."""
            if rank <= 0:                       # rank-0 baseline: no low-rank branch (pure W4A8)
                return T.new_zeros(T.shape[0], 0), T.new_zeros(0, T.shape[1])
            if M is not None:
                B = T @ M                                        # ‖T−L‖ in act-metric = ‖(T−L)M‖
                U, S, Vh = torch.linalg.svd(B, full_matrices=False)
                r_ = min(rank, S.shape[0])
                return U[:, :r_] * S[:r_].view(1, -1), Vh[:r_, :] @ Minv
            try:
                U, S, Vh = torch.linalg.svd(T, full_matrices=False)
            except RuntimeError:                                  # OOM -> CPU
                U, S, Vh = torch.linalg.svd(T.cpu(), full_matrices=False)
                U, S, Vh = U.to(T.device), S.to(T.device), Vh.to(T.device)
            r_ = min(rank, S.shape[0])
            return U[:, :r_] * S[:r_].view(1, -1), Vh[:r_, :]

        def recon_err(up, down, Rq):
            """Reconstruction error of (L+Q(R)) vs Ŵ, in the metric the fit optimizes."""
            E = Wd - up @ down - Rq
            if M is not None:
                return float(torch.linalg.matrix_norm(E @ M))
            return float(torch.linalg.matrix_norm(E))

        # 2+3. iterative low-rank fit + 4-bit residual with error feedback (keep best).
        up, down = lowrank_fit(Wd)
        Rq = _quant_weight(Wd - up @ down, w_fmt, w_bits, w_group)
        best = (up, down, Rq); best_err = recon_err(up, down, Rq); best_it = 0
        for it in range(1, refine_iters + 1):
            up, down = lowrank_fit(Wd - Rq)                      # absorb the quant error into L
            Rq = _quant_weight(Wd - up @ down, w_fmt, w_bits, w_group)
            e = recon_err(up, down, Rq)
            if e < best_err:
                best_err, best, best_it = e, (up, down, Rq), it
        up, down, Rq = best
        lora_up, lora_down, W_res_fake = up.to(dev), down.to(dev), Rq.to(dev)
        r = lora_up.shape[1]

        m = cls(in_f, out_f, lin.bias is not None, r, w_bits, a_bits, w_group,
                a_group=a_group, dtype=dtype, w_fmt=w_fmt, a_fmt=a_fmt)
        m = m.to(dev)
        m.smooth.copy_(s.to(dtype))
        m.lora_down.copy_(lora_down.to(dtype))
        m.lora_up.copy_(lora_up.to(dtype))
        m.w_res.copy_(W_res_fake.to(dtype))
        if lin.bias is not None:
            m.bias.copy_(lin.bias.data.to(dtype))

        # diagnostic: plain weight-recon rel-err (comparable across modes); if whitening, also
        # the activation-weighted output rel-err the method actually optimizes.
        What_approx = lora_up @ lora_down + W_res_fake
        W_approx = What_approx / s.view(1, -1)
        rel = float((W_approx - W).norm() / (W.norm() + 1e-8))
        diag = {"rel_err": rel, "rank": r, "in": in_f, "out": out_f,
                "whiten": bool(use_whiten), "refine_best_it": best_it}
        if use_whiten:
            E = (What - What_approx).to(svd_device)
            out_err = torch.sqrt(torch.trace(E @ Ghat @ E.t()).clamp(min=0))
            out_den = torch.sqrt(torch.trace(Wd @ Ghat @ Wd.t()).clamp(min=0))
            diag["out_rel_err"] = float(out_err / (out_den + 1e-8))
        return m, diag


# --------------------------------------------------------------------------------------
# module-tree surgery helpers
# --------------------------------------------------------------------------------------
def _get_parent(root, dotted):
    parent = root
    parts = dotted.split(".")
    for p in parts[:-1]:
        parent = parent[int(p)] if p.isdigit() else getattr(parent, p)
    return parent, parts[-1]


def _set_module(root, dotted, mod):
    parent, leaf = _get_parent(root, dotted)
    if leaf.isdigit():
        parent[int(leaf)] = mod
    else:
        setattr(parent, leaf, mod)


def target_linear_names(transformer, prefixes=("transformer_blocks.", "single_transformer_blocks.")):
    """All nn.Linear module names under the given prefixes (the block attn/MLP projections)."""
    names = []
    for name, mod in transformer.named_modules():
        if isinstance(mod, nn.Linear) and any(name.startswith(p) for p in prefixes):
            names.append(name)
    return names


def collect_act_stats(transformer, target_names, with_gram=True, gram_device="cuda"):
    """Forward-pre-hooks that accumulate, per target Linear, BOTH:
      - per-input-channel running max-abs (for the smoothing factor), and
      - the raw activation Gram  G = Σₜ xₜ xₜᵀ  (in×in, for activation-aware/whitened SVD).

    The Gram is accumulated in fp32 on `gram_device` (≈18 GB total for klein-4B on an
    80 GB A100). One calibration pass populates everything (smoothing is diagonal, so the
    smoothed Gram derives from the raw one at decompose time — no second pass).

    Returns (absmax dict, gram dict|None, handles). Remove handles, then read.
    """
    absmax = {n: None for n in target_names}
    gram = {n: None for n in target_names} if with_gram else None
    target = set(target_names)
    handles = []

    def mk(n):
        def hook(mod, inp):
            x = inp[0].detach()
            xf = x.reshape(-1, x.shape[-1]).float()                 # (T, in)
            a = xf.abs().amax(dim=0).cpu()
            absmax[n] = a if absmax[n] is None else torch.maximum(absmax[n], a)
            if with_gram:
                g = (xf.to(gram_device).t() @ xf.to(gram_device))  # (in, in)
                gram[n] = g if gram[n] is None else gram[n] + g
        return hook

    for name, mod in transformer.named_modules():
        if name in target:
            handles.append(mod.register_forward_pre_hook(mk(name)))
    return absmax, gram, handles


def collect_act_absmax(transformer, target_names):
    """Back-compat shim: absmax-only collection (plain SVD path)."""
    absmax, _, handles = collect_act_stats(transformer, target_names, with_gram=False)
    return absmax, handles


@torch.no_grad()
def apply_svdquant_from_stats(transformer, stats, rank=32, alpha=0.5, w_bits=4, a_bits=8,
                              w_group=64, a_group=0, svd_device="cuda", grams=None,
                              whiten=True, gram_ridge=1e-3, refine_iters=3, smooth=True,
                              w_fmt="int", a_fmt="int"):
    """Replace every target Linear (those with collected stats) by an SVDQuantLinear.

    `stats` = per-layer activation abs-max; `grams` = per-layer raw activation Gram XᵀX
    (required for whiten=True activation-aware SVD; if None, falls back to plain SVD).

    Returns (specs, diags): `specs` is the per-layer shape/rank record needed to rebuild
    empty modules before load_state_dict; `diags` is per-layer recon rel-err (weight, and
    output-rel-err when whitening).
    """
    specs, diags = {}, {}
    for name, absmax in stats.items():
        if absmax is None:
            continue
        parent, leaf = _get_parent(transformer, name)
        lin = parent[int(leaf)] if leaf.isdigit() else getattr(parent, leaf)
        g = grams.get(name) if (grams is not None) else None
        m, diag = SVDQuantLinear.from_linear(lin, absmax, rank=rank, alpha=alpha,
                                             w_bits=w_bits, a_bits=a_bits, w_group=w_group,
                                             a_group=a_group, svd_device=svd_device,
                                             act_gram=g, whiten=whiten, gram_ridge=gram_ridge,
                                             refine_iters=refine_iters, smooth=smooth,
                                             w_fmt=w_fmt, a_fmt=a_fmt)
        _set_module(transformer, name, m)
        specs[name] = {"in": m.in_features, "out": m.out_features,
                       "bias": m.bias is not None, "rank": m.rank}
        diags[name] = diag
        if g is not None:
            grams[name] = None  # free the Gram as we go (memory)
    return specs, diags


def apply_svdquant_empty(transformer, specs, w_bits=4, a_bits=8, w_group=64, a_group=0,
                         dtype=None, w_fmt="int", a_fmt="int"):
    """Rebuild empty SVDQuantLinear modules from saved specs (for load_state_dict)."""
    if dtype is None:
        dtype = next(transformer.parameters()).dtype
    dev = next(transformer.parameters()).device
    for name, sp in specs.items():
        m = SVDQuantLinear(sp["in"], sp["out"], sp["bias"], sp["rank"],
                           w_bits=w_bits, a_bits=a_bits, w_group=w_group, a_group=a_group,
                           dtype=dtype, w_fmt=w_fmt, a_fmt=a_fmt)
        _set_module(transformer, name, m.to(dev))
    return transformer


def quant_summary(transformer):
    """Effective compressed size + count of quantized layers."""
    qbytes = full_bytes = 0
    nq = 0
    for _, m in transformer.named_modules():
        if isinstance(m, SVDQuantLinear):
            qbytes += m.compressed_bytes()
            full_bytes += m.out_features * m.in_features * 2  # bf16 baseline
            nq += 1
    return {"n_quant_layers": nq, "quant_MB": qbytes / 1e6, "full_MB": full_bytes / 1e6,
            "ratio": (full_bytes / qbytes) if qbytes else 0.0}