diff options
| -rw-r--r-- | ep_run/casc_eq_train.py | 69 |
1 files changed, 52 insertions, 17 deletions
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index f360df1..dfb7098 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -62,6 +62,11 @@ ap.add_argument('--adam_b1', type=float, default=0.9) # AdamW beta1 (late-SN ap.add_argument('--est', choices=['single', 'centered', 'richardson'], default='single') # centered: [g(+b)+g(-b)]/2 (O(b^2) bias, 2x relax cost) # richardson: 2g(b)-g(2b) (O(b^2) bias, large-b friendly) +ap.add_argument('--est_late', choices=['', 'centered'], default='') +ap.add_argument('--est_late_at', type=int, default=0) # switch --est -> --est_late at this step (process-local, + # bf_late_at semantics); centered is TAIL medicine +ap.add_argument('--centfast', action='store_true') # centered via ONE doubled batch [x;x], +beta/-beta halves + # (shared kernels; math identical to sequential centered) args = ap.parse_args() if args.olmo2: args.untie = True @@ -300,11 +305,13 @@ def tok_sigma(iters=8): v = W.t() @ u; sig = v.norm(); v /= max(sig, 1e-12) return float(sig) -def relax(z0, zs, ins, outs, y, beta, K, x): +def relax(z0, zs, ins, outs, y, beta, K, x, bmask=None): """K fb rounds with GRAPH REUSE + two dedups: (a) the top CE force d_top is refreshed on even rounds only (states move O(beta) per round -> O(beta^2) error); (b) the LAST rebuild 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.""" + reuses them instead of re-running a full graphed chain. + bmask (rows,1,1): per-row multiplier on the top force (centfast +/-1 halves); the + rows*T-aware scale keeps per-row d identical to the sequential B-sized run.""" d = [None] * args.L geta_l = args.geta adaptive = args.relax_tol > 0 @@ -313,7 +320,10 @@ def relax(z0, zs, ins, outs, y, beta, K, x): 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() + nbt_loc = zc.shape[0] * zc.shape[1] + g = torch.autograd.grad(ce, zc)[0] + if bmask is not None: g = g * bmask + d[args.L - 1] = (-beta * nbt_loc * g).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() @@ -408,24 +418,38 @@ def ep_step(x, y): beta_t = beta_t * GOV.get('cap', 1.0) # wall-2 loop-gain cap OVERRIDES the floor (the ceiling # can sit below the floor near the wall; survival first) if args.bsign_rand and torch.rand((), generator=BGEN).item() < 0.5: beta_t = -beta_t - z0, zs, ins, outs = free_states_graphed(x) + EST = args.est + if args.est_late and GOV['step'] >= args.est_late_at: EST = args.est_late + CF = (EST == 'centered' and args.centfast) + if CF: # doubled batch [x;x]: +beta half / -beta half share every kernel (holofast pattern) + x_in, y_in = torch.cat([x, x], 0), torch.cat([y, y], 0) + bmask = torch.ones(x_in.shape[0], 1, 1, device=dev); bmask[args.B:] = -1.0 + halves = (slice(0, args.B), slice(args.B, None)) + else: + x_in, y_in, bmask, halves = x, y, None, (slice(None),) + z0, zs, ins, outs = free_states_graphed(x_in) 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, last_outs = relax(z0, zs, ins, outs, y, +beta_t, GOV['K'], x) + free_ce = F.cross_entropy(readout(zs_free[-1][:args.B]).reshape(-1, vocab), y.reshape(-1)).item() + zp, last_outs = relax(z0, zs, ins, outs, y_in, +beta_t, GOV['K'], x_in, bmask=bmask) + def _drift(zp_, zf_): + dr = 0.0 + for h in halves: # per-half worst drift == sequential guard decisions (max over passes) + num = sum(float((a[h] - b[h]).norm()) for a, b in zip(zp_, zf_)) + den = sum(float(b[h].norm()) for b in zf_) + dr = max(dr, num / max(den, 1e-9)) + return dr 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) + drift = _drift(zp, zs_free) gdrift = ddp_max_scalar(drift) # guard DECISIONS on the global worst -> identical on every rank if (not math.isfinite(gdrift)) or (gdrift > 0.5 and not args.noguard): ok_retry = False if args.kretry > 0 and math.isfinite(gdrift) and not args.noguard: GOV['skr'] = GOV.get('skr', 0) + 1 # marginal batch: retry once with deeper relaxation - z0, zs, ins, outs = free_states_graphed(x) + z0, zs, ins, outs = free_states_graphed(x_in) zs_free = [z.clone() for z in zs] - zp, last_outs = relax(z0, zs, ins, outs, y, +beta_t, args.kretry, x) + zp, last_outs = relax(z0, zs, ins, outs, y_in, +beta_t, args.kretry, x_in, bmask=bmask) 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) + drift = _drift(zp, zs_free) gdrift = ddp_max_scalar(drift) ok_retry = math.isfinite(gdrift) and gdrift <= 0.5 if not ok_retry: @@ -442,19 +466,30 @@ def ep_step(x, y): GOV['cap'] = max(GOV.get('cap', 1.0) * 0.85, 0.05) # attack (gentler than v1) elif res_g < 0.01 or rho_g < 0.5 * args.beta_cap_rho: GOV['cap'] = min(GOV.get('cap', 1.0) * 1.02, 1.0) # recover - 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) - if args.est == 'single': + if CF: + # one-graph centered: [g(+b)+g(-b)]/2 = d[(E+ - E-)/(2b·NBT)]/dtheta; CE-head term from + # the +beta half only (matches sequential centered's gsC at the +beta top states). + Ec = 0.0 + for z, o in zip(zp, last_outs): + df = z.detach().float() - o.float() + Ec = Ec + 0.5 * (df[:args.B] ** 2).sum() - 0.5 * (df[args.B:] ** 2).sum() + obj = Ec / (NBT * 2.0 * beta_t) + obj_loss(readout(zp[-1][:args.B].detach()).reshape(-1, vocab), y.reshape(-1)) + gs = torch.autograd.grad(obj, all_params, allow_unused=True) + elif EST == 'single': + 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) obj = E / (NBT * beta_t) + obj_loss(readout(zp[-1].detach()).reshape(-1, vocab), y.reshape(-1)) gs = torch.autograd.grad(obj, all_params, allow_unused=True) else: # two-pass estimators: g(b) := d[E(b)]/dtheta / (NBT*b) => single-sided bias g_true + c*b. # centered: [g(+b) + g(-b)] / 2 (1/b sign inside => average cancels c*b) # richardson: 2*g(b) - g(2b) (extrapolation cancels c*b at large b) + E = 0.0 + for z, o in zip(zp, last_outs): E = E + 0.5 * ((z.detach().float() - o.float()) ** 2).sum() gsE = torch.autograd.grad(E / (NBT * beta_t), all_params, allow_unused=True) gsC = torch.autograd.grad(obj_loss(readout(zp[-1].detach()).reshape(-1, vocab), y.reshape(-1)), all_params, allow_unused=True) - b2 = -beta_t if args.est == 'centered' else 2.0 * beta_t + b2 = -beta_t if EST == 'centered' else 2.0 * beta_t z0b, zsb, insb, outsb = free_states_graphed(x) zsb_free = [z.clone() for z in zsb] zpb, lob = relax(z0b, zsb, insb, outsb, y, b2, GOV['K'], x) @@ -473,7 +508,7 @@ def ep_step(x, y): if a is None and b is None: return None a = a if a is not None else torch.zeros_like(b) b = b if b is not None else torch.zeros_like(a) - return (a + b) / 2.0 if args.est == 'centered' else (2.0 * a - b) + return (a + b) / 2.0 if EST == 'centered' else (2.0 * a - b) gs = [(_comb(e, e2) if (e is not None or e2 is not None) else None) for e, e2 in zip(gsE, gsE2)] gs = [ (g if g is not None else c) if c is None or g is None else g + c for g, c in zip(gs, gsC) ] gs = ddp_avg(gs, all_params) # global-batch gradient; gn/gema/guard below see identical values on all ranks |