1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
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()
|
