diff options
| -rw-r--r-- | ep_run/casc_eq_train.py | 56 |
1 files changed, 48 insertions, 8 deletions
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index 60d2297..5214311 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -51,6 +51,9 @@ ap.add_argument('--data', default='tinystories_bpe') # dataset dir under ep_ ap.add_argument('--sync_check', type=int, default=500) # DDP: verify bitwise param sync every N steps (0=off) ap.add_argument('--ddp_backend', default='nccl', choices=['nccl', 'gloo']) # gloo = correctness tests on shared GPUs ap.add_argument('--ddp_grad_test', action='store_true') # one-step grad equivalence test vs single-GPU big batch, then exit +ap.add_argument('--relax_tol', type=float, default=0.0) # >0: ADAPTIVE relax — sweep until rel. state change < tol + # (or --kmax), geta backtracks x0.6 on residual GROWTH (rho>=1 + # signal), then one final graphed round. 0 = legacy fixed-K. ap.add_argument('--muon_mom', type=float, default=0.95) # Muon momentum (late-SNR arm: 0.99 = ~10x noise averaging) ap.add_argument('--adam_b1', type=float, default=0.9) # AdamW beta1 (late-SNR arm companion) ap.add_argument('--est', choices=['single', 'centered', 'richardson'], default='single') @@ -300,16 +303,25 @@ def relax(z0, zs, ins, outs, y, beta, K, x): keeps graphs (layer-0 fed a graphed emb) and returns (ins, outs) so the theta-readout reuses them instead of re-running a full graphed chain.""" d = [None] * args.L - for k in range(K): - if k % args.dtop_every == 0 or d[args.L - 1] is None: + geta_l = args.geta + adaptive = args.relax_tol > 0 + + def forces(refresh_top): + if refresh_top or d[args.L - 1] is None: zc = zs[args.L - 1].detach().requires_grad_(True) 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].to(outs[l + 1].dtype))[0].detach().float() - last = (k + 1 == K) + + def rebuild(last): + nonlocal ins, outs prev = z0 n_ins, n_outs = [], [] + rnum = rden = 0.0 + g_eff = 1.0 if last else geta_l # FINAL graphed round is ALWAYS full-step: the theta-read + # identity (z - o) = d requires undamped substitution; + # mixing there leaks the iteration residual into E (gn 1e5 bug) with torch.autocast('cuda', dtype=torch.bfloat16, enabled=args.amp): for l in range(args.L): if last and l == 0: @@ -317,10 +329,38 @@ def relax(z0, zs, ins, outs, y, beta, K, x): else: i = prev.detach().requires_grad_(True) o = blocks[l](i, mask) - zs[l] = (o.detach().float() + d[l]) + znew = o.detach().float() + d[l] + # damped (under-relaxed) mixing: geta<1 restores contraction on stiff operators + # (wall-2 toolkit); fixed point unchanged (z = z + geta*(o+d-z) <=> z = o+d) + mixed = znew if g_eff >= 1.0 else (zs[l] + g_eff * (znew - zs[l])) + with torch.no_grad(): + rnum += float((mixed - zs[l]).norm()); rden += float(zs[l].norm()) + zs[l] = mixed n_ins.append(i); n_outs.append(o) prev = zs[l] ins, outs = n_ins, n_outs + return rnum / max(rden, 1e-9) + + if not adaptive: # legacy fixed-K path (bit-identical to before) + for k in range(K): + forces(k % args.dtop_every == 0) + rebuild(k + 1 == K) + GOV['kuse'] = K + return zs, outs + + prev_res, k = None, 0 + while k < args.kmax: + forces(k % args.dtop_every == 0) + res = rebuild(False) + k += 1 + if prev_res is not None and res > prev_res and res > args.relax_tol: + geta_l = max(0.2, geta_l * 0.6) # residual GREW: local rho>=1 -> damp harder + prev_res = res + if res < args.relax_tol: + break + forces(True) # final graphed round at the settled state (theta-read) + rebuild(True) + GOV['kuse'] = k + 1 return zs, outs def dFdtheta(zs, x, y, beta): @@ -380,7 +420,7 @@ def ep_step(x, y): if not ok_retry: GOV['skd'] = GOV.get('skd', 0) + 1 # drift-guard reject (relaxation non-convergence) for p in all_params: p.grad = None - return free_ce, beta_t, GOV['K'], False + return free_ce, beta_t, GOV.get('kuse', GOV['K']), False GOV['drift'] = gdrift E = 0.0 for z, o in zip(zp, last_outs): E = E + 0.5 * ((z.detach().float() - o.float()) ** 2).sum() # fp32 accumulation (bf16-safe; no-op in fp32) @@ -405,7 +445,7 @@ def ep_step(x, y): if (not math.isfinite(gdrift2)) or (gdrift2 > 0.5 and not args.noguard): GOV['skd'] = GOV.get('skd', 0) + 1 # second-pass drift reject -> skip step (synced) for p in all_params: p.grad = None - return free_ce, beta_t, GOV['K'], False + return free_ce, beta_t, GOV.get('kuse', GOV['K']), False E2 = 0.0 for z, o in zip(zpb, lob): E2 = E2 + 0.5 * ((z.detach().float() - o.float()) ** 2).sum() gsE2 = torch.autograd.grad(E2 / (NBT * b2), all_params, allow_unused=True) @@ -427,10 +467,10 @@ def ep_step(x, y): if not math.isfinite(gn) or (gn > 8 * GOV['gema'] and not args.noguard): GOV['skg'] = GOV.get('skg', 0) + 1 # gn-EMA-guard reject (gradient-magnitude spike) for p in all_params: p.grad = None - return free_ce, beta_t, GOV['K'], False + return free_ce, beta_t, GOV.get('kuse', GOV['K']), False for p, g in zip(all_params, gs): p.grad = g - return free_ce, beta_t, GOV['K'], True + return free_ce, beta_t, GOV.get('kuse', GOV['K']), True def bp_gate(x, y): """true BP grads for telemetry cos (called before opt.step; reads p.grad separately).""" |