diff options
| -rw-r--r-- | ep_run/casc_eq_train.py | 70 |
1 files changed, 47 insertions, 23 deletions
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index 547312e..392d2be 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -22,6 +22,8 @@ ap.add_argument('--kmax', type=int, default=8) # adaptive fb rounds cap ap.add_argument('--noguard', action='store_true') # diagnosis: skip only non-finite grads ap.add_argument('--untie', action='store_true') # separate readout matrix (untied from tok) ap.add_argument('--tok_init', type=float, default=0.0) # >0: init tok/pos with this std (GPT-standard 0.02) +ap.add_argument('--compile', action='store_true') # torch.compile each block (free speed where supported) +ap.add_argument('--sig_every', type=int, default=25) # tok-sigma refresh interval (amortized) ap.add_argument('--gate_every', type=int, default=200) # in-training cos(EP,BP) telemetry args = ap.parse_args() torch.manual_seed(args.seed) @@ -53,6 +55,12 @@ if args.tok_init > 0: with torch.no_grad(): tok.weight.normal_(0, args.tok_init); pos.weight.normal_(0, args.tok_init) blocks = nn.ModuleList([Block(args.C, args.H) for _ in range(args.L)]).to(dev) +if args.compile: + try: + for i in range(args.L): blocks[i] = torch.compile(blocks[i], mode='reduce-overhead') + print('[compile] blocks compiled', flush=True) + except Exception as e: + print(f'[compile] disabled ({e})', flush=True) mask = torch.triu(torch.full((args.T, args.T), float('-inf'), device=dev), 1) W_out = nn.Parameter(torch.randn(vocab, args.C, device=dev) * 0.02) if args.untie else None readout = (lambda z: z @ W_out.t()) if args.untie else (lambda z: z @ tok.weight.t()) @@ -61,13 +69,18 @@ opt = torch.optim.AdamW(all_params, lr=args.lr, weight_decay=1e-4) sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1.0, (s + 1) / max(args.warmup, 1))) NBT = args.B * args.T -def free_states(x): +def free_states_graphed(x): + """free forward, keeping per-layer graphs (in_l, out_l) so round-1 backward vjps reuse them.""" with torch.no_grad(): - z = tok(x) + pos(torch.arange(args.T, device=dev))[None] - z0 = z.clone(); zs = [] - for b in blocks: - z = b(z, mask); zs.append(z) - return z0, zs + z0 = (tok(x) + pos(torch.arange(args.T, device=dev))[None]) + ins, outs, zs = [], [], [] + prev = z0 + for b in blocks: + i = prev.detach().requires_grad_(True) + o = b(i, mask) + ins.append(i); outs.append(o); zs.append(o.detach()) + prev = zs[-1] + return z0, zs, ins, outs @torch.no_grad() def tok_sigma(iters=8): @@ -80,25 +93,33 @@ def tok_sigma(iters=8): v = W.t() @ u; sig = v.norm(); v /= max(sig, 1e-12) return float(sig) -def relax(z0, zs_free, y, beta, K): - """K fb rounds: backward feedback refresh + forward rebuild. State oscillation is - HARMLESS for the theta-readout (probe-verified); no contraction verdict.""" - zs = [z.clone() for z in zs_free] +def relax(z0, zs, ins, outs, y, beta, K): + """K fb rounds with GRAPH REUSE: the backward vjps consume the graphs stored by the + previous forward (free pass for round 1, rebuild pass afterwards) — saves one full + graphed chain per round. Rebuild keeps graphs for the next round; oscillation harmless.""" d = [None] * args.L for k in range(K): zc = zs[args.L - 1].detach().requires_grad_(True) ce = F.cross_entropy(readout(zc).reshape(-1, vocab), y.reshape(-1)) d[args.L - 1] = (-beta * NBT * torch.autograd.grad(ce, zc)[0]).detach() for l in range(args.L - 2, -1, -1): - zc = zs[l].detach().requires_grad_(True) - fnext = blocks[l + 1](zc, mask) - d[l] = torch.autograd.grad(fnext, zc, grad_outputs=d[l + 1])[0].detach() - with torch.no_grad(): - prev = z0 - for l in range(args.L): - rebuilt = blocks[l](prev, mask) + d[l] - zs[l] = (1 - args.geta) * zs[l] + args.geta * rebuilt if args.geta < 1.0 else rebuilt - prev = zs[l] + # d_l = J_{l+1}^T d_{l+1}; outs[l+1] was computed at input ins[l+1] == zs[l] + d[l] = torch.autograd.grad(outs[l + 1], ins[l + 1], grad_outputs=d[l + 1], + retain_graph=(k + 1 < K and False))[0].detach() + last = (k + 1 == K) + prev = z0 + n_ins, n_outs = [], [] + for l in range(args.L): + i = prev.detach().requires_grad_(True) + if last: + with torch.no_grad(): + o = blocks[l](i, mask) + else: + o = blocks[l](i, mask) + zs[l] = (o.detach() + d[l]) + n_ins.append(i); n_outs.append(o) + prev = zs[l] + ins, outs = n_ins, n_outs return zs def dFdtheta(zs, x, y, beta): @@ -117,13 +138,16 @@ def ep_step(x, y): K = GOV['K'] fb rounds; guard = finiteness + drift + grad-norm sanity only.""" global SIG0 if GOV['K'] is None: GOV['K'] = args.K - sig = tok_sigma() - GOV['sig'] = sig + if GOV.get('step', 0) % args.sig_every == 0 or GOV.get('sig', 0) == 0: + GOV['sig'] = tok_sigma() + GOV['step'] = GOV.get('step', 0) + 1 + sig = GOV['sig'] if SIG0 is None: SIG0 = sig beta_t = args.beta * GOV['bscale'] * (SIG0 * SIG0) / max(sig * sig, 1e-9) - z0, zs_free = free_states(x) + z0, zs, ins, outs = free_states_graphed(x) + zs_free = [z.clone() for z in zs] free_ce = F.cross_entropy(readout(zs_free[-1]).reshape(-1, vocab), y.reshape(-1)).item() - zp = relax(z0, zs_free, y, +beta_t, GOV['K']) + zp = relax(z0, zs, ins, outs, y, +beta_t, GOV['K']) with torch.no_grad(): drift = sum(float((a - b).norm()) for a, b in zip(zp, zs_free)) / max( sum(float(b.norm()) for b in zs_free), 1e-9) |