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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn

2 files changed, 44 insertions, 6 deletions

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)