summaryrefslogtreecommitdiff
path: root/ep_run/bp_lm.py
diff options
context:
space:
mode:
authorYuren Hao <yurenh2@illinois.edu>2026-07-06 02:22:54 -0500
committerYuren Hao <yurenh2@illinois.edu>2026-07-06 02:22:54 -0500
commit8f43de671cf03ee20d5d652b6cd3ba575982a436 (patch)
tree590ab1b4ccb249d0fb15205d2480e69b7bfba080 /ep_run/bp_lm.py
parent66b4ad978585d1b8008e02d06787811e2ee93da7 (diff)
bp_lm twin control + env forensics probes + holo_a_track_avg (semi-convergence fix, ungated)
BP+EMA 2x2 (local, healthy env): qknorm {1.9951, 1.9977}, no-qknorm {1.9728, 1.9732} -> parameterization-matched BP twin tops at ~1.97 vs EP 1.7888: the equilibrium computation's iteration/depth dividend = 0.18 CE from identical parameters. External anchor (tuned depth-1 BP, 1.7921): EP at parity. bp_lm on 107/2.3.1 gave 1.9823 ~= local -> plain backward exonerated on the pascal env; the 107 divergence (pair/floss/resreg dead by step 600) narrows to the EP reg/estimator loop. Single-step fingerprints all match (a/b/c/d) -> suspected intermittent kernel issue; discriminators in flight (torch-2.7 env probe + delay_hr2's step-2000 reg-on transition). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn
Diffstat (limited to 'ep_run/bp_lm.py')
-rw-r--r--ep_run/bp_lm.py75
1 files changed, 75 insertions, 0 deletions
diff --git a/ep_run/bp_lm.py b/ep_run/bp_lm.py
new file mode 100644
index 0000000..7fe0bdd
--- /dev/null
+++ b/ep_run/bp_lm.py
@@ -0,0 +1,75 @@
+"""bp_lm.py — the plain-BP standard-transformer reference with weight-EMA (task #34), rebuilt for
+metric parity with the EP lineage: SAME parameter tensors as EQBlock (7.48M), SAME data + eval
+protocol, and the SAME pema weight-EMA so both lineages report best-of(raw, ema) val. Forward = one
+standard pre-LN residual block (no relaxation): h = x + attn(LN1 x) + ffn(LN2 x); logits = h @ Wh.
+Historical anchor (no EMA, no qknorm): 'BP standard transformer (7.48M, lr 3e-3) best 1.7921 (20k)'
+(FINDINGS:355). Run with --qknorm for architecture parity with the current EP recipe."""
+import argparse, time, torch, torch.nn.functional as F
+import lt_ep_train as L
+
+
+def fwd(blk, idx):
+ x = blk.embed(idx)
+ h1 = F.layer_norm(x, (blk.C,), blk.ln1g, blk.ln1b)
+ h2 = F.layer_norm(x, (blk.C,), blk.ln2g, blk.ln2b)
+ h = x + blk.attn(h1) + (F.gelu(h2 @ blk.fc + blk.fcb, approximate='tanh') @ blk.pj + blk.pjb)
+ return h @ blk.Wh
+
+
+def evaluate(blk, nb=8, B=32):
+ tot = 0.0
+ with torch.no_grad():
+ for _ in range(nb):
+ idx, y = L.get_batch('val', B, blk.T)
+ tot += float(F.cross_entropy(fwd(blk, idx).reshape(-1, L.vocab), y.reshape(-1)))
+ return tot / nb
+
+
+def main():
+ ap = argparse.ArgumentParser()
+ ap.add_argument('--steps', type=int, default=32000)
+ ap.add_argument('--B', type=int, default=24)
+ ap.add_argument('--lr', type=float, default=3e-3)
+ ap.add_argument('--wd', type=float, default=1e-4)
+ ap.add_argument('--pema', type=float, default=0.999)
+ ap.add_argument('--qknorm', action='store_true')
+ ap.add_argument('--log', type=int, default=200)
+ ap.add_argument('--ckpt', type=str, default='runs/bp_lm.pt')
+ cfg = ap.parse_args()
+ torch.manual_seed(0)
+ blk = L.EQBlock(512, 16, 256, 256, c=1.0, attn_mode='thick')
+ blk.qknorm = cfg.qknorm
+ opt = torch.optim.AdamW(blk.allp, lr=cfg.lr, weight_decay=cfg.wd)
+ sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, cfg.steps, eta_min=cfg.lr * 0.05)
+ ema = [p.detach().clone() for p in blk.allp]
+ best, t0 = float('inf'), time.time()
+ for step in range(1, cfg.steps + 1):
+ idx, y = L.get_batch('train', cfg.B, blk.T)
+ loss = F.cross_entropy(fwd(blk, idx).reshape(-1, L.vocab), y.reshape(-1))
+ opt.zero_grad(set_to_none=True)
+ loss.backward()
+ opt.step(); sched.step()
+ with torch.no_grad():
+ for e, p in zip(ema, blk.allp):
+ e.mul_(cfg.pema).add_(p, alpha=1 - cfg.pema)
+ if step % cfg.log == 0:
+ raw = evaluate(blk)
+ bak = [p.detach().clone() for p in blk.allp]
+ with torch.no_grad():
+ for p, e in zip(blk.allp, ema):
+ p.copy_(e)
+ emav = evaluate(blk)
+ with torch.no_grad():
+ for p, b in zip(blk.allp, bak):
+ p.copy_(b)
+ m = min(raw, emav)
+ if m < best:
+ best = m
+ torch.save({'allp': [p.detach().cpu() for p in blk.allp], 'step': step, 'best': best}, cfg.ckpt)
+ print(f"step {step:5d}/{cfg.steps} | val {raw:.4f} ema {emav:.4f} (best {best:.4f}) "
+ f"| {step / (time.time() - t0):.2f} it/s", flush=True)
+ print(f"[bp_lm] DONE best {best:.4f}", flush=True)
+
+
+if __name__ == '__main__':
+ main()