| """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 |
|
|
|
|
| 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) |
| 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/<name>/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): |
| 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) |
|
|
