"""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())