From 643983aec870af4ba7e71eafc6c1ba37b2348730 Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Mon, 13 Jul 2026 23:27:01 -0500 Subject: Estimator arms machinery: --est centered/richardson (two-pass, O(b^2) bias), --muon_mom/--adam_b1 knobs; smoke passed (resume+centered cos 0.998) Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn --- ep_run/casc_eq_train.py | 40 ++++++++++++++++++++++++++++++++++++++-- ep_run/muon.py | 10 ++++++---- 2 files changed, 44 insertions(+), 6 deletions(-) (limited to 'ep_run') diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index 18b58e1..60d2297 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -51,6 +51,11 @@ 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('--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') + # 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) args = ap.parse_args() if args.olmo2: args.untie = True @@ -235,6 +240,7 @@ if args.bf16: if args.opt == 'muon': from muon import build_hybrid opt, sched = build_hybrid(blocks, all_params, args.lr, args.muon_lr, args.warmup, + muon_mom=args.muon_mom, adam_b1=args.adam_b1, total_steps=(args.steps if args.cosine else 0), lr_min_ratio=args.lr_min_ratio) else: if args.wd >= 0: # OLMo2-style grouped decay: linear weights + head decay; embeddings/norm-gains none @@ -378,8 +384,38 @@ def ep_step(x, y): 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) - 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) + if args.est == 'single': + 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) + 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 + 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) + with torch.no_grad(): + drift2 = sum(float((a - b).norm()) for a, b in zip(zpb, zsb_free)) / max( + sum(float(b.norm()) for b in zsb_free), 1e-9) + gdrift2 = ddp_max_scalar(drift2) + 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 + 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) + def _comb(a, b): + 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) + 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 gn = 0.0 for g in gs: diff --git a/ep_run/muon.py b/ep_run/muon.py index 7281811..d65a231 100644 --- a/ep_run/muon.py +++ b/ep_run/muon.py @@ -56,15 +56,17 @@ class MultiSched: for s in self.scheds: s.step() -def build_hybrid(blocks, other_params, lr_adamw, lr_muon, warmup, total_steps=0, lr_min_ratio=0.1): +def build_hybrid(blocks, other_params, lr_adamw, lr_muon, warmup, total_steps=0, lr_min_ratio=0.1, + muon_mom=0.95, adam_b1=0.9): """Muon(2D block matrices) + AdamW(everything else). Scheds: linear warmup, then cosine decay to - lr_min_ratio*peak if total_steps>0 (long runs), else constant after warmup (legacy).""" + lr_min_ratio*peak if total_steps>0 (long runs), else constant after warmup (legacy). + muon_mom/adam_b1: momentum knobs (late-SNR noise-averaging arms, 2026-07-13).""" import math as _m mats = [p for p in blocks.parameters() if p.ndim == 2] mat_ids = {id(p) for p in mats} rest = [p for p in other_params if id(p) not in mat_ids] - om = Muon(mats, lr=lr_muon) - oa = torch.optim.AdamW(rest, lr=lr_adamw, weight_decay=1e-4) + om = Muon(mats, lr=lr_muon, momentum=muon_mom) + oa = torch.optim.AdamW(rest, lr=lr_adamw, weight_decay=1e-4, betas=(adam_b1, 0.999)) if total_steps > 0: def fn(s): if s < warmup: return (s + 1) / max(warmup, 1) -- cgit v1.2.3