diff options
Diffstat (limited to 'ep_run/casc_eq_train.py')
| -rw-r--r-- | ep_run/casc_eq_train.py | 33 |
1 files changed, 18 insertions, 15 deletions
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index 91a1cca..6a05cf5 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -42,6 +42,7 @@ ap.add_argument('--bf_late', type=float, default=0.0) # >0: raise beta_floor ap.add_argument('--bf_late_at', type=int, default=25000) ap.add_argument('--bsign_rand', action='store_true') # random-sign beta per step (KHS 'random scheme'): averages the O(beta) single-sided bias at single-phase cost ap.add_argument('--bf16', action='store_true') # cast model to bf16 (E-accumulation + tok_sigma stay fp32) — the x0.5 cost lever, GATE before production +ap.add_argument('--amp', action='store_true') # PROPER mixed precision: autocast(bf16) matmuls, fp32 params/states/d/E — amp_gate.py PASSED 2026-07-12 (cos 0.9682 vs fp32 0.9687); --bf16 naive-cast stays DEAD (state quantization, RESULT 11) ap.add_argument('--dtop_every', type=int, default=1) # 1 = exact (DEFAULT, BP-parity); 2 = fast mode (~20% cheaper, ~4% CE tax at high lr) ap.add_argument('--gate_every', type=int, default=200) # in-training cos(EP,BP) telemetry; <=0 = fully BP-free (no bp_gate at all) ap.add_argument('--gate_govern', action='store_true') # let gate cos adjust K/bscale (default: observe-only => training control is BP-free) @@ -219,11 +220,12 @@ def free_states_graphed(x): z0 = emb(x) 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] + with torch.autocast('cuda', dtype=torch.bfloat16, enabled=args.amp): + for b in blocks: + i = prev.detach().requires_grad_(True) + o = b(i, mask) + ins.append(i); outs.append(o); zs.append(o.detach().float()) + prev = zs[-1] return z0, zs, ins, outs @torch.no_grad() @@ -249,19 +251,20 @@ def relax(z0, zs, ins, outs, y, beta, K, x): ce = obj_loss(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): - d[l] = torch.autograd.grad(outs[l + 1], ins[l + 1], grad_outputs=d[l + 1])[0].detach() + d[l] = torch.autograd.grad(outs[l + 1], ins[l + 1], grad_outputs=d[l + 1].to(outs[l + 1].dtype))[0].detach().float() last = (k + 1 == K) prev = z0 n_ins, n_outs = [], [] - for l in range(args.L): - if last and l == 0: - i = emb(x) # graphed emb for the readout's E-path - else: - i = prev.detach().requires_grad_(True) - o = blocks[l](i, mask) - zs[l] = (o.detach() + d[l]) - n_ins.append(i); n_outs.append(o) - prev = zs[l] + with torch.autocast('cuda', dtype=torch.bfloat16, enabled=args.amp): + for l in range(args.L): + if last and l == 0: + i = emb(x) # graphed emb for the readout's E-path + else: + i = prev.detach().requires_grad_(True) + o = blocks[l](i, mask) + zs[l] = (o.detach().float() + d[l]) + n_ins.append(i); n_outs.append(o) + prev = zs[l] ins, outs = n_ins, n_outs return zs, outs |