testbed/hyper_expert.py

"""HyperExpert -- a per-token GENERATED low-rank expert, drop-in for `Bank`.

Same forward interface as Bank: takes the MLP input hidden state x and returns
the MLP output y of the same shape. Instead of a fixed E-column bank, a small
hypernetwork generates a per-token rank-r FFN from the token's own hidden state,
on top of a shared per-layer low-width base FFN.

For hidden dim d:
  encoder : z = gelu(Enc(x)),  Enc: Linear(d -> c),     z in R^c
  generator: U = (G_U z).reshape(r, d),   G_U in R^{(r*d) x c}
             V = (G_V z).reshape(d, r),   G_V in R^{(d*r) x c}
  base    : a stored per-layer width-b FFN, base(x) = Wb2 gelu(Wb1 x + bb1) + bb2
  output  : y = base(x) + V @ gelu(U @ x)

Per-layer stored footprint ~ 2*d*r*c (generators) + 2*d*b (base) + d*c (encoder).

CHUNKED generation (`chunk=N`, the DECODE-SPEED lever).  The generator is the
expensive per-token cost; with N>1 we fire it ONCE per N-token chunk and reuse
the generated (U,V) across the chunk, amortizing the generator ~N-fold.  The
shared base FFN STILL RUNS PER TOKEN, so per-token processing is preserved --
only the generated low-rank correction V@gelu(U@x) is held constant within a
chunk.  chunk=1 reduces exactly to the per-token path above (sanity contract).

CHUNK MODES (`chunk_mode`).  "fixed" (default) = the original arange//N chunking.
"sentence" = boundary-aware chunking: a chunk never crosses a sentence/line
boundary ( . ! ? ; newline) and holds at most `chunk_cap` tokens; a long segment
is greedily split into <=cap sub-chunks (still coherent).  Per-token chunk ids are
computed tokenizer-side by chunk_ids_from_tokens() and handed in via set_chunk_ids()
(prefill/train) or via set_stream_boundary() per step (streaming decode).  Same rule
both regimes: one generated (U,V) per chunk, base FFN per token; differentiable.

Two forward regimes share one rule (one (U,V) per chunk, base per token):
  * PREFILL / TRAIN (`_stream=False`, x is [B,S,d]): vectorized.  Each chunk's
    latent z is the FULL-CHUNK MEAN of that chunk's inputs.  NOTE: full-chunk
    mean peeks at within-chunk future tokens, so this is the *prefill-style*
    pooling -- documented and fine for the matched-budget ppl comparison (the
    held-out eval is a full forward).  A strictly-causal variant is the prefix/
    running mean; at a chunk boundary in a stream that prefix is just the first
    token, which is exactly what the decode path below uses.
  * DECODE / STREAM (`_stream=True`, called one token at a time by HF.generate):
    regenerate (U,V) only when crossing a chunk boundary (global pos % N == 0),
    else reuse the cached (U,V).  This yields exactly ceil(tokens/N) generator
    calls -> the real amortization measured by bench_chunk.py.  Boundary pooling
    is the first token of the chunk (the causal prefix available at that step).
    The tiny prefill(mean) vs decode(first-token) pooling difference only shifts
    generated-text content; the reported numbers are ppl (prefill) and tok/s
    (decode), so neither is corrupted.  Call reset_stream() before each stream.

Warm start: G_V is zero-initialised so the generated correction starts at exactly
0 -> at step 0 the expert IS a width-b bank (y == base(x), equal to Bank modulo
compute dtype; verified relRMSE 2.4e-3 in diag_warmstart.py). It then learns the
per-token low-rank correction on top of that bank. This is the real warm start:
the hypernet begins EXACTLY where the bank begins.

Base subset (init): EMPIRICALLY, init="random" (default) is the right choice.
The diag_warmstart.py sweep (30-layer simultaneous replacement, fixed held-out)
measured INIT held-out ppl:
    bank_random_fp32  11874   bank_topnorm_fp32  42502
    hyper_random_bf16 12490   hyper_topnorm_bf16 42493
i.e. the baseline Bank itself uses RANDOM init (train_compress2.py default) and
starts at ~11.9k -- NOT "a few hundred". Top-norm is ~3.5x WORSE here: across 30
stacked layers the dropped-neuron errors of a top-norm subset compound coherently.
So init="random" both matches the bank baseline and gives the lower INIT. (g_v=0
makes the hyper start ~equal to whichever bank you pick; random is the good bank.)
init="topnorm" is kept only for the bisection that established the above.
"""
import torch, torch.nn as nn, torch.nn.functional as F

# ---- SENTENCE/LINE-boundary chunking helpers --------------------------------
# A "chunk" never crosses a sentence/line boundary AND is at most `cap` tokens.
# Boundaries on the CODE testbed = newlines + sentence-enders ( . ! ? ; \n ).
# A long segment (run of tokens up to and including its boundary ender) is split
# into sub-chunks of <= cap tokens (greedy: start a new chunk at a boundary token
# OR when the open chunk reaches `cap`).  These two functions are tokenizer-side
# (they map token ids -> per-token chunk ids); the HyperExpert itself consumes the
# resulting integer chunk-id tensor and never needs the tokenizer.

BOUNDARY_CHARS = ".!?;\n"


def boundary_token_mask(tok, chars=BOUNDARY_CHARS):
    """Bool tensor [vocab]: True if a token's decoded text contains a boundary char.
    Such a token ENDS its segment (the next token must begin a new chunk)."""
    size = len(tok)
    mask = torch.zeros(size, dtype=torch.bool)
    cset = set(chars)
    for i in range(size):
        txt = tok.decode([i])
        if any(ch in cset for ch in txt):
            mask[i] = True
    return mask


def chunk_ids_from_tokens(ids, bmask, cap):
    """Per-token chunk ids [B,S] (long) under the boundary+cap rule.
    Scan each row: token s starts a new chunk iff the PREVIOUS token was a
    boundary (segment ended) OR the open chunk already holds `cap` tokens. So a
    chunk is contiguous, <= cap tokens, and a boundary token is always its LAST
    token -> chunks never cross a boundary.  cids are constants (no grad)."""
    B, S = ids.shape
    bnd = bmask.to("cpu")[ids.to("cpu")].tolist()      # [B][S] python bools
    out = [[0] * S for _ in range(B)]
    for b in range(B):
        row, orow = bnd[b], out[b]
        cur = 0; length = 0
        for s in range(S):
            if s == 0:
                cur = 0; length = 1
            elif row[s - 1] or length >= cap:
                cur += 1; length = 1
            else:
                length += 1
            orow[s] = cur
    return torch.tensor(out, dtype=torch.long, device=ids.device)


class HyperExpert(nn.Module):
    def __init__(self, src_mlp, c, r, b, dtype=torch.bfloat16, init="random", chunk=1,
                 chunk_mode="fixed", chunk_cap=20):
        super().__init__()
        d = src_mlp.c_fc.weight.shape[1]          # in/out dim (3072)
        self.d, self.c, self.r, self.b = d, c, r, b
        self.chunk = max(1, int(chunk))           # tokens per generated expert (1=per-token)
        # chunk_mode: "fixed" = arange//chunk (original); "sentence" = boundary+cap
        # chunking driven by per-token chunk ids fed in from the tokenizer side.
        self.chunk_mode = chunk_mode
        self.chunk_cap = max(1, int(chunk_cap))   # max tokens/chunk in sentence mode (N)
        self.dtype = dtype                        # param/compute dtype (bf16 keeps

        # encoder: x -> z in R^c
        self.enc = nn.Linear(d, c)

        # generators: z -> flattened U (r*d) and V (d*r)
        self.g_u = nn.Linear(c, r * d, bias=False)
        self.g_v = nn.Linear(c, d * r, bias=False)

        # shared per-layer width-b base FFN, warm-started from the SAME top-norm
        # b-column subset of the original MLP that Bank(init="topnorm") uses, so
        # base(x) == Bank(x) at init. (init="random" reproduces the old cold start.)
        if init == "topnorm":
            idx = src_mlp.c_proj.weight.data.float().norm(dim=0).topk(b).indices
        else:
            idx = torch.randperm(src_mlp.c_fc.weight.shape[0])[:b]
        self.b_fc = nn.Linear(d, b)
        self.b_proj = nn.Linear(b, d)
        with torch.no_grad():
            self.b_fc.weight.copy_(src_mlp.c_fc.weight.data[idx].float())
            self.b_fc.bias.copy_(src_mlp.c_fc.bias.data[idx].float()
                                 if src_mlp.c_fc.bias is not None else torch.zeros(b))
            self.b_proj.weight.copy_(src_mlp.c_proj.weight.data[:, idx].float())
            self.b_proj.bias.copy_(src_mlp.c_proj.bias.data.float()
                                   if src_mlp.c_proj.bias is not None else torch.zeros(d))

        # small encoder/G_U init; zero G_V so initial correction is 0 (y == base)
        nn.init.normal_(self.enc.weight, std=0.02); nn.init.zeros_(self.enc.bias)
        nn.init.normal_(self.g_u.weight, std=0.02)
        nn.init.zeros_(self.g_v.weight)

        # Store params in `dtype` (bf16 by default): the hypernetwork is ~800M
        # params/expert; fp32 params+grads for all 30 experts would peak near the
        # 16 GB GPU limit and OOM at the first backward. bf16 halves that footprint.
        self.to(dtype)
        self.last_in = None; self.last_out = None
        # sentence-mode chunk ids for the CURRENT vectorized/prefill forward [B,S],
        # set from outside (train/eval/prompt) via set_chunk_ids(); never grad.
        self._chunk_ids = None
        # streaming (decode) state -- see reset_stream(); only used when _stream=True
        self._stream = False
        self.reset_stream()

    # sentence-mode: per-token chunk ids for the next vectorized forward (prefill).
    def set_chunk_ids(self, cids):
        self._chunk_ids = cids; return self

    # sentence-mode streaming: boundary flag(s) for the incoming token(s) [B].
    def set_stream_boundary(self, bnd):
        self._stream_bnd = bnd; return self

    # ---- streaming (autoregressive decode) controls --------------------------
    def set_streaming(self, flag):
        self._stream = bool(flag); return self

    def reset_stream(self):
        """Reset the per-generation streaming state. Call before each decode."""
        self._pos = 0          # global token position in the current stream
        self._Uc = None        # cached generated U for the active chunk
        self._Vc = None        # cached generated V for the active chunk
        self._calls = 0        # generator firings so far (for amortization stats)
        # sentence-mode streaming extras (batch=1 decode, as bench uses):
        self._open_len = 0     # tokens accumulated in the current open chunk
        self._pending = False  # last token was a boundary -> next token opens a chunk
        self._stream_bnd = None  # boundary flag(s) for the incoming token [B]
        self._prefill_calls = 0  # generator firings during the prompt prefill

    def base(self, x):
        return self.b_proj(F.gelu(self.b_fc(x), approximate="tanh"))

    # ---- chunked correction paths -------------------------------------------
    def _chunk_prefill(self, xs):
        """Vectorized full-chunk-mean correction. xs: [B,S,d] -> res: [B,S,d]."""
        B, S, d = xs.shape; N = self.chunk
        n_chunks = (S + N - 1) // N
        cidx = torch.arange(S, device=xs.device) // N         # [S] chunk id per token
        pooled = xs.new_zeros(B, n_chunks, d)
        pooled.index_add_(1, cidx, xs)                        # sum inputs per chunk
        counts = torch.bincount(cidx, minlength=n_chunks).to(xs.dtype).clamp_min(1)
        pooled = pooled / counts.view(1, n_chunks, 1)         # full-chunk mean -> [B,nc,d]
        z = F.gelu(self.enc(pooled), approximate="tanh")      # [B,nc,c]   ONE z per chunk
        U = self.g_u(z).view(B, n_chunks, self.r, d)          # [B,nc,r,d] ONE expert per chunk
        V = self.g_v(z).view(B, n_chunks, d, self.r)          # [B,nc,d,r]
        Ut = U.index_select(1, cidx)                          # [B,S,r,d] broadcast to tokens
        Vt = V.index_select(1, cidx)                          # [B,S,d,r]
        h = F.gelu(torch.matmul(Ut, xs.unsqueeze(-1)).squeeze(-1), approximate="tanh")  # [B,S,r]
        return torch.matmul(Vt, h.unsqueeze(-1)).squeeze(-1)  # [B,S,d]

    def _gen(self, x1):
        """Fire the generator once from a single pooled vector x1 [B,d]; cache (U,V)."""
        B = x1.shape[0]
        z = F.gelu(self.enc(x1), approximate="tanh")          # [B,c]
        self._Uc = self.g_u(z).view(B, self.r, self.d)        # [B,r,d]
        self._Vc = self.g_v(z).view(B, self.d, self.r)        # [B,d,r]
        self._calls += 1

    def _chunk_stream(self, xs):
        """Stateful decode correction. xs: [B,S,d] -> res: [B,S,d].
        Regenerates (U,V) only at chunk boundaries; reuses the cache otherwise."""
        B, S, d = xs.shape; N = self.chunk
        if S > 1:
            # one-shot prompt prefill: same full-chunk-mean correction, counted as
            # ceil(S/N) generator firings; seed the cache from the trailing token so
            # subsequent S==1 decode steps continue cleanly.
            res = self._chunk_prefill(xs)
            self._calls += (S + N - 1) // N
            self._gen(xs[:, -1, :]); self._calls -= 1  # cache only, not an extra count
            self._pos += S
            return res
        x1 = xs[:, 0, :]                                       # [B,d] single decode token
        if (self._pos % N == 0) or (self._Uc is None):
            self._gen(x1)                                     # boundary -> fire generator
        h = F.gelu(torch.bmm(self._Uc, x1.unsqueeze(-1)).squeeze(-1), approximate="tanh")  # [B,r]
        res = torch.bmm(self._Vc, h.unsqueeze(-1)).squeeze(-1)                              # [B,d]
        self._pos += 1
        return res.unsqueeze(1)                                # [B,1,d]

    # ---- SENTENCE-mode chunked correction paths -----------------------------
    def _chunk_prefill_ids(self, xs, cids):
        """Vectorized chunk-mean correction with PER-ROW chunk ids (sentence mode).
        xs: [B,S,d], cids: [B,S] long -> res: [B,S,d]. One generated (U,V) per
        chunk (chunk mean of its inputs), gathered back to every token. Differentiable
        w.r.t. xs (scatter_add/gather/matmul); cids are constants."""
        B, S, d = xs.shape
        NC = int(cids.max().item()) + 1
        idx = cids.unsqueeze(-1).expand(B, S, d)              # [B,S,d]
        pooled = xs.new_zeros(B, NC, d).scatter_add_(1, idx, xs)         # sum per chunk
        counts = xs.new_zeros(B, NC).scatter_add_(1, cids, xs.new_ones(B, S))
        pooled = pooled / counts.clamp_min(1).unsqueeze(-1)  # chunk mean -> [B,NC,d]
        z = F.gelu(self.enc(pooled), approximate="tanh")     # [B,NC,c] ONE z per chunk
        U = self.g_u(z).view(B, NC, self.r, d)               # [B,NC,r,d]
        V = self.g_v(z).view(B, NC, d, self.r)               # [B,NC,d,r]
        gu = cids.view(B, S, 1, 1).expand(B, S, self.r, d)
        gv = cids.view(B, S, 1, 1).expand(B, S, d, self.r)
        Ut = U.gather(1, gu)                                 # [B,S,r,d] per-token expert
        Vt = V.gather(1, gv)                                 # [B,S,d,r]
        h = F.gelu(torch.matmul(Ut, xs.unsqueeze(-1)).squeeze(-1), approximate="tanh")  # [B,S,r]
        return torch.matmul(Vt, h.unsqueeze(-1)).squeeze(-1)  # [B,S,d]

    def _chunk_stream_sentence(self, xs):
        """Stateful decode correction for sentence mode (batch=1). Regenerates (U,V)
        when the previous token was a boundary OR the open chunk hit `chunk_cap`."""
        B, S, d = xs.shape; cap = self.chunk_cap
        if S > 1:
            # prompt prefill: full chunk-mean correction over the prompt's own chunks.
            cids = self._chunk_ids
            res = self._chunk_prefill_ids(xs, cids)
            nc = int(cids.max().item()) + 1
            self._calls += nc; self._prefill_calls = nc
            last = int(cids[0, -1].item())
            self._open_len = int((cids[0] == last).sum().item())  # size of trailing chunk
            self._gen(xs[:, -1, :]); self._calls -= 1             # seed cache, not a count
            bnd = self._stream_bnd
            self._pending = bool(bnd.any().item()) if bnd is not None else False
            self._pos += S
            return res
        x1 = xs[:, 0, :]                                       # [B,d] single decode token
        start_new = (self._Uc is None) or self._pending or (self._open_len >= cap)
        if start_new:
            self._gen(x1); self._open_len = 1                 # boundary/cap -> fire generator
        else:
            self._open_len += 1
        bnd = self._stream_bnd
        self._pending = bool(bnd.any().item()) if bnd is not None else False
        h = F.gelu(torch.bmm(self._Uc, x1.unsqueeze(-1)).squeeze(-1), approximate="tanh")  # [B,r]
        res = torch.bmm(self._Vc, h.unsqueeze(-1)).squeeze(-1)                              # [B,d]
        self._pos += 1
        return res.unsqueeze(1)                                # [B,1,d]

    def forward(self, x):
        self.last_in = x
        shape = x.shape
        d = self.d
        if self.chunk_mode == "sentence":
            # ---- SENTENCE/LINE-boundary + capped path (one (U,V) per chunk) ----
            if x.ndim == 3:
                B, S = shape[0], shape[1]
            else:
                B, S = 1, x.reshape(-1, d).shape[0]
            xs = x.reshape(B, S, d).to(self.dtype)
            if self._stream:
                res = self._chunk_stream_sentence(xs)
            else:
                cids = self._chunk_ids
                assert cids is not None and tuple(cids.shape) == (B, S), \
                    "sentence mode prefill needs set_chunk_ids([B,S]) before forward"
                res = self._chunk_prefill_ids(xs, cids)
            y = self.base(xs) + res                            # base PER TOKEN + chunk corr.
            out = y.view(shape).to(x.dtype)
            self.last_out = out
            return out
        if self.chunk <= 1:
            # ---- per-token path (original; --chunk 1 reproduces this exactly) ----
            x2 = x.reshape(-1, d).to(self.dtype)               # [N, d]
            N = x2.shape[0]
            z = F.gelu(self.enc(x2), approximate="tanh")       # [N, c]
            U = self.g_u(z).view(N, self.r, d)                 # [N, r, d]
            V = self.g_v(z).view(N, d, self.r)                 # [N, d, r]
            h = F.gelu(torch.bmm(U, x2.unsqueeze(-1)).squeeze(-1), approximate="tanh")  # [N, r]
            res = torch.bmm(V, h.unsqueeze(-1)).squeeze(-1)    # [N, d]
            y = self.base(x2) + res                            # [N, d]
            if self._stream: self._calls += N                  # per-token: one firing/token
            out = y.view(shape).to(x.dtype)
            self.last_out = out
            return out
        # ---- chunked path (one generated expert per N tokens; base per token) ----
        if x.ndim == 3:
            B, S = shape[0], shape[1]
        else:
            B, S = 1, x.reshape(-1, d).shape[0]
        xs = x.reshape(B, S, d).to(self.dtype)                 # [B, S, d]
        res = self._chunk_stream(xs) if self._stream else self._chunk_prefill(xs)
        y = self.base(xs) + res                                # base PER TOKEN + chunked corr.
        out = y.view(shape).to(x.dtype)
        self.last_out = out
        return out

    def footprint(self):
        return sum(p.numel() for p in self.parameters())