"""Post-hoc V2 mask alignment on LFM2.5-230M-Base (local MPS pilot). Student = LFM with sparse MLPs, trained by distillation from the frozen original (teacher). Student MLP: g = w1 x, u = w3 x; mask = top-K of |g| (K = FRAC of the hidden size, fixed for the whole run); h = mask * silu(g) * u (values kept, not binarised). Layers in --dense stay dense. V2 loss (as nanochat_patch/gpt.py v2_score, sign-split): units = [+ halves | - halves] of the neurons; forward mask = hard signed mask, backward through sigmoid(+g) / sigmoid(-g) (straight-through). Target T = relu(cos of batch-centred MLP outputs of the *teacher* at the same layer and positions), so T cannot collapse. Same V2_TOKENS random positions for every layer per micro-batch (--micro rows; --bs rows per optimizer step via gradient accumulation). Loss = KL(teacher || student) per token + lam * mean over sparse layers of (1 - v2). Data: documents from a ClimbMix parquet shard (--data), BOS-aligned rows: every row starts with BOS and a new document. Eval: bpb on lfm_topk_probe.texts(EVAL_DOCS) (held-out shard; disjoint from training docs when --data is a train shard). python lfm_posthoc.py --name lam0p1 --lam 0.1 --tokens 20e6 [--data val_text/shard_00000.parquet] """ import argparse, json, math, os, time import numpy as np, pyarrow.parquet as pq, torch, torch.nn.functional as F from huggingface_hub import hf_hub_download from transformers import AutoModelForCausalLM, PreTrainedTokenizerFast import lfm_topk_probe as P from contextlib import nullcontext ap = argparse.ArgumentParser() ap.add_argument("--name", required=True); ap.add_argument("--lam", type=float, default=0.0) ap.add_argument("--tokens", type=float, default=20e6); ap.add_argument("--seq", type=int, default=1024) ap.add_argument("--bs", type=int, default=8); ap.add_argument("--micro", type=int, default=2); ap.add_argument("--lr", type=float, default=5e-5) ap.add_argument("--warmup", type=int, default=100); ap.add_argument("--frac", type=float, default=0.25) ap.add_argument("--dense", default="0,1"); ap.add_argument("--v2-tokens", type=int, default=1024) ap.add_argument("--data", default=os.path.join(P.ROOT, "val_text", "shard_00000.parquet")) ap.add_argument("--doc-start", type=int, default=0); ap.add_argument("--eval-every", type=int, default=500) ap.add_argument("--eval-docs", type=int, default=64); ap.add_argument("--max-steps", type=int, default=0) args = ap.parse_args() if __name__ == "__main__" else None dev = P.dev; torch.manual_seed(0) amp = (lambda: torch.autocast("cuda", torch.bfloat16)) if dev == "cuda" else nullcontext # bf16 compute, fp32 weights def v2_score(m, y): """Uncentred cosine of masks vs relu(cosine of batch-centred outputs), off-diagonal pairs (nanochat_patch/gpt.py).""" m, y = m.float(), y.float() Km = F.normalize(m, dim=-1, eps=1e-8); Km = Km @ Km.T Ky = F.normalize(y - y.mean(0, keepdim=True), dim=-1, eps=1e-8); Ky = torch.relu(Ky @ Ky.T) off = 1.0 - torch.eye(m.size(0), device=m.device); Km, Ky = Km * off, Ky * off return (Km * Ky).sum() / (Km.norm() * Ky.norm() + 1e-8) class Sparse: """Patches the student's MLPs. self.idx = sampled flat positions for V2 (None: no V2); self.target[l] = teacher y.""" def __init__(self, model, frac, dense): self.idx, self.target, self.v2 = None, {}, {} self.sparse = [l for l in range(len(model.model.layers)) if l not in dense] for l in self.sparse: m = model.model.layers[l].feed_forward; m.forward = self.make(l, m, int(frac * m.w1.out_features)) def make(self, l, m, K): def fwd(x): g = m.w1(x); u = m.w3(x); s = F.silu(g) hard = torch.zeros_like(g).scatter_(-1, g.abs().topk(K, -1).indices, 1.0) y = m.w2(hard * s * u) if self.idx is not None: gi = g.flatten(0, 1)[self.idx]; hi = hard.flatten(0, 1)[self.idx] hp, hn = hi * (gi > 0), hi * (gi < 0); sp, sn = torch.sigmoid(gi), torch.sigmoid(-gi) ms = torch.cat([hp + sp - sp.detach(), hn + sn - sn.detach()], -1) # forward hard, backward sigmoid self.v2[l] = v2_score(ms, self.target[l]) return y return fwd def rows(tok, path, seq, start): """BOS-aligned rows of seq+1 tokens: each row starts with BOS + a new document; documents are cut at the row end.""" texts = pq.read_table(path, columns=["text"]).column("text").to_pylist()[start:] bos = tok.bos_token_id; buf = [] for t in texts: ids = tok(t, add_special_tokens=False).input_ids if not buf: buf = [bos] else: buf += [bos] buf += ids if len(buf) >= seq + 1: yield buf[:seq + 1]; buf = [] @torch.no_grad() def eval_bpb(model, tok, docs, nbytes): tot = 0.0 for t in docs: x = torch.cat([torch.tensor([[tok.bos_token_id]]), tok(t, return_tensors="pt", add_special_tokens=False).input_ids], 1).to(dev) with amp(): lg = model(x).logits tot += F.cross_entropy(lg[0, :-1].float(), x[0, 1:], reduction="sum").item() return tot / math.log(2) / nbytes def load_student(name, frac=0.25, dense=(0, 1)): """Trained student from nanochat_runs/lfm_posthoc//student.pt, sparse MLPs patched (V2 off).""" m = AutoModelForCausalLM.from_pretrained(P.MODEL, dtype=torch.float32) m.load_state_dict(torch.load(os.path.join(P.RUNS, "lfm_posthoc", name, "student.pt"), map_location="cpu")) m = m.to(P.dev).eval(); Sparse(m, frac, set(dense)); return m if __name__ == "__main__": OUT = os.path.join(P.RUNS, "lfm_posthoc", args.name); os.makedirs(OUT, exist_ok=True) tok = PreTrainedTokenizerFast(tokenizer_file=hf_hub_download(P.MODEL, "tokenizer.json"), bos_token="<|startoftext|>") teacher = AutoModelForCausalLM.from_pretrained(P.MODEL, dtype=torch.bfloat16).to(dev).eval().requires_grad_(False) student = AutoModelForCausalLM.from_pretrained(P.MODEL, dtype=torch.float32).to(dev).train() S = Sparse(student, args.frac, {int(x) for x in args.dense.split(",") if x != ""}) tgt = {} def thook(l): def f(mod, a, o): if S.idx is not None: tgt[l] = o.flatten(0, 1)[S.idx].detach() return f for l in S.sparse: teacher.model.layers[l].feed_forward.register_forward_hook(thook(l)) docs = P.texts(args.eval_docs); nbytes = sum(len(t.encode()) for t in docs) opt = torch.optim.AdamW(student.parameters(), lr=args.lr, betas=(0.9, 0.95), weight_decay=0.0) steps = args.max_steps or int(args.tokens // (args.bs * args.seq)) sched = lambda s: min(1, (s + 1) / args.warmup) * 0.5 * (1 + math.cos(math.pi * min(1, s / steps))) log = []; student.eval(); b0 = eval_bpb(student, tok, docs, nbytes); student.train() print(f"{args.name}: {steps} steps x {args.bs * args.seq} tokens, lam {args.lam}, sparse layers {S.sparse}; " f"bpb at start {b0:.4f}", flush=True); log.append(dict(step=0, bpb=b0)) it = rows(tok, args.data, args.seq, args.doc_start); t0 = time.time() for step in range(1, steps + 1): for gp in opt.param_groups: gp["lr"] = args.lr * sched(step) opt.zero_grad(set_to_none=True); n_acc = args.bs // args.micro; kl_s = v2_s = 0.0 for _ in range(n_acc): # micro-batches: logits are B x T x 65536 x = torch.tensor([next(it) for _ in range(args.micro)], device=dev)[:, :-1] S.idx = torch.randint(0, x.numel(), (args.v2_tokens,), device=dev) with torch.no_grad(): tl = F.log_softmax(teacher(x).logits.float(), -1) S.target = tgt with amp(): sl = student(x).logits sl = F.log_softmax(sl.float(), -1) kl = F.kl_div(sl, tl, log_target=True, reduction="none").sum(-1).mean(); del tl v2 = torch.stack([S.v2[l] for l in S.sparse]).mean() loss = kl + args.lam * (1 - v2) if args.lam > 0 else kl (loss / n_acc).backward(); kl_s += kl.item() / n_acc; v2_s += v2.item() / n_acc torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0); opt.step() if step % 25 == 0 or step == 1: el = time.time() - t0 print(f"step {step}/{steps} kl {kl_s:.4f} v2 {v2_s:.4f} | {step * args.bs * args.seq / el:.0f} tok/s, " f"eta {(steps - step) * el / step / 60:.0f} min", flush=True) log.append(dict(step=step, kl=kl_s, v2=v2_s)) if step % args.eval_every == 0 or step == steps: S.idx = None; student.eval(); b = eval_bpb(student, tok, docs, nbytes); student.train() print(f" eval step {step}: bpb {b:.4f}", flush=True); log.append(dict(step=step, bpb=b)) json.dump(dict(args=vars(args), log=log), open(os.path.join(OUT, "log.json"), "w"), indent=1) S.idx = None torch.save(student.state_dict(), os.path.join(OUT, "student.pt")) print("saved ->", OUT)