From 40be67d4f5b5a6b46c662c70b759e585e136d70e Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Mon, 6 Jul 2026 09:42:52 -0500 Subject: Tier-3 gates: Anderson math-yes/impl-no (parked for v2); bf16polish UNSAFE near-edge (eval-only); dp_ep.py ready MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Anderson: res 25-35x deeper per budget but 5.8x slower (naive history stacks + per-iter safeguard eval) — v2 = ring buffers + periodic safeguard, est +1.3x on the speed tier. bf16+20polish: 1.41x free phase, res parity, BUT z-diff 1.2e-3 — near-marginal operators contract too slowly for a 20-step polish (0.998^20≈0.96), same magnitude as the TF32 kill verdict and the estimator's 50%-sensitivity input. Predicted by our own depth/noise theory. Flags kept with warnings; neither ships for training. dp_ep.py: manual-allreduce EP DP (controller in lockstep, aligned collectives), smoke pending freed 1080s. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn --- ep_run/dp_ep.py | 71 ++++++++++++++++++++++++++++++++++++++++++++++++++ ep_run/fastfp_gate.log | 7 +++++ ep_run/fastfp_gate.py | 40 ++++++++++++++++++++++++++++ ep_run/lt_ep_train.py | 46 ++++++++++++++++++++++++++++++++ 4 files changed, 164 insertions(+) create mode 100644 ep_run/dp_ep.py create mode 100644 ep_run/fastfp_gate.log create mode 100644 ep_run/fastfp_gate.py (limited to 'ep_run') diff --git a/ep_run/dp_ep.py b/ep_run/dp_ep.py new file mode 100644 index 0000000..0d76911 --- /dev/null +++ b/ep_run/dp_ep.py @@ -0,0 +1,71 @@ +"""dp_ep.py — data-parallel EP (task #14's last item). EP has no autograd backward, so DDP does not +apply; this is the minimal manual-allreduce loop: identical init on every rank (same seed), per-rank +batches, ep_step -> all_reduce(mean) over the canonical param order (zeros for absent grads so the +collective stays aligned), all_reduce the residual so the jr controller stays in lockstep, identical +AdamW steps -> weights never diverge. Launch: + torchrun --nproc_per_node=N dp_ep.py --steps 300 [recipe flags...] +Effective batch = B * world (lr scaling deliberately NOT applied for the smoke; tune later).""" +import argparse, os, time, torch +import torch.distributed as dist +import lt_ep_train as L + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument('--steps', type=int, default=300) + ap.add_argument('--B', type=int, default=24) + ap.add_argument('--lr', type=float, default=6e-4) + ap.add_argument('--seed', type=int, default=0) + ap.add_argument('--reg_delay', type=int, default=10 ** 9) # smoke default: reg-free (Pascal-safe) + ap.add_argument('--resreg', type=float, default=0.2) + ap.add_argument('--jacreg', type=float, default=0.1) + ap.add_argument('--t2sel', type=int, default=40) + ap.add_argument('--hr', type=float, default=0.02) + ap.add_argument('--holofast', action='store_true') + ap.add_argument('--sdpa', action='store_true') + ap.add_argument('--log', type=int, default=50) + cfg = ap.parse_args() + + dist.init_process_group('nccl') + rank, world = dist.get_rank(), dist.get_world_size() + torch.cuda.set_device(rank) + L.dev = f'cuda:{rank}' + + torch.manual_seed(cfg.seed) # identical init on every rank + blk = L.EQBlock(512, 16, 256, 256, c=1.0, attn_mode='thick') + blk.qknorm = True; blk.track = True + blk.holofast, blk.sdpa = cfg.holofast, cfg.sdpa + for p in blk.allp: + dist.broadcast(p.data, 0) # belt & suspenders: exact bitwise identical start + torch.manual_seed(cfg.seed * 1009 + rank + 1) # per-rank data stream + opt = torch.optim.AdamW(blk.allp, lr=cfg.lr, weight_decay=1e-4) + + jr = cfg.jacreg + t0 = time.time() + for step in range(1, cfg.steps + 1): + idx, y = L.get_batch('train', cfg.B, blk.T) + dly = step < cfg.reg_delay + grads, res = L.ep_step(blk, idx, y, 150, 20, 0.1, 0.02, 0.0 if dly else jr, holo=2, hr=cfg.hr, + t1max=300, res_est=1e-4, t2sel=cfg.t2sel, corr_every=1, res_gate=0.0, + resreg=0.0 if dly else cfg.resreg) + rt = torch.tensor([res], device=L.dev) + dist.all_reduce(rt); res = float(rt) / world # controller signal identical across ranks + for p in blk.allp: # canonical order: aligned collectives + g = grads.get(id(p)) + if g is None: + g = torch.zeros_like(p) + dist.all_reduce(g) + g /= world + p.grad = g + opt.step(); opt.zero_grad(set_to_none=True) + if step % cfg.log == 0 and rank == 0: + v = L.evaluate(blk, 150, 0.1, nb=2) + print(f"[dp{world}] step {step}/{cfg.steps} | val {v:.4f} | res {res:.1e} | " + f"{step / (time.time() - t0):.3f} it/s(x{world}B)", flush=True) + if rank == 0: + print(f"[dp{world}] DONE {cfg.steps} steps in {time.time() - t0:.0f}s", flush=True) + dist.destroy_process_group() + + +if __name__ == '__main__': + main() diff --git a/ep_run/fastfp_gate.log b/ep_run/fastfp_gate.log new file mode 100644 index 0000000..5beee9a --- /dev/null +++ b/ep_run/fastfp_gate.log @@ -0,0 +1,7 @@ +[s2000] euler150 : 1.85s res=9.98e+00 +[s2000] anderson : 10.75s res=3.88e-01 evals=150 z-diff=5.01e-02 +[s2000] bf16+20 : 1.31s res=9.96e+00 z-diff=1.19e-03 +[fast-adaptive] euler150 : 1.64s res=3.31e+00 +[fast-adaptive] anderson : 10.54s res=9.49e-02 evals=150 z-diff=3.81e-02 +[fast-adaptive] bf16+20 : 1.21s res=3.33e+00 z-diff=1.13e-03 +FASTFP_GATE_DONE diff --git a/ep_run/fastfp_gate.py b/ep_run/fastfp_gate.py new file mode 100644 index 0000000..f65f71f --- /dev/null +++ b/ep_run/fastfp_gate.py @@ -0,0 +1,40 @@ +"""Gate for --fastfp (Anderson z*) and --bf16polish: endpoint parity + residual + eval-count/time vs the +150-step Euler reference, on two operators (deeply-trained fast-adaptive + near-edge s2000 — the hard +case: Anderson historically fails on cycling ops, s2000's marginal band is the stress test). +GPU shared with abl_pair — timings are relative, parity/evals exact.""" +import time, torch +import lt_ep_train as L + +for name, path in (('s2000', 'runs/redx_traj/s2000.pt'), ('fast-adaptive', 'runs/ep_fast_adaptive.pt')): + torch.manual_seed(0) + blk = L.EQBlock(512, 16, 256, 256, c=1.0, attn_mode='thick'); blk.qknorm = True + ck = torch.load(path, map_location=L.dev) + with torch.no_grad(): + for p, w in zip(blk.allp, ck['allp']): + p.copy_(w.to(L.dev)) + torch.manual_seed(42) + idx, _ = L.get_batch('train', 24, 256) + xin = blk.embed(idx).detach() + + torch.cuda.synchronize(); t = time.time() + z_ref = L.relax(blk, xin.clone(), xin, 150, 0.1) + torch.cuda.synchronize(); t_ref = time.time() - t + r_ref = (L.relax(blk, z_ref, xin, 1, 0.1) - z_ref).norm().item() + print(f"[{name}] euler150 : {t_ref:5.2f}s res={r_ref:.2e}", flush=True) + + torch.cuda.synchronize(); t = time.time() + z_aa, evals = L.anderson_relax(blk, xin.clone(), xin, 150, 0.1) + torch.cuda.synchronize(); t_aa = time.time() - t + r_aa = (L.relax(blk, z_aa, xin, 1, 0.1) - z_aa).norm().item() + zd = ((z_aa - z_ref).norm() / (z_ref.norm() + 1e-12)).item() + print(f"[{name}] anderson : {t_aa:5.2f}s res={r_aa:.2e} evals={evals} z-diff={zd:.2e}", flush=True) + + blk.bf16polish = 20 + torch.cuda.synchronize(); t = time.time() + z_bf = L.relax(blk, xin.clone(), xin, 150, 0.1) + torch.cuda.synchronize(); t_bf = time.time() - t + blk.bf16polish = 0 + r_bf = (L.relax(blk, z_bf, xin, 1, 0.1) - z_bf).norm().item() + zdb = ((z_bf - z_ref).norm() / (z_ref.norm() + 1e-12)).item() + print(f"[{name}] bf16+20 : {t_bf:5.2f}s res={r_bf:.2e} z-diff={zdb:.2e}", flush=True) +print("FASTFP_GATE_DONE", flush=True) diff --git a/ep_run/lt_ep_train.py b/ep_run/lt_ep_train.py index 76bb9bb..0fd9180 100644 --- a/ep_run/lt_ep_train.py +++ b/ep_run/lt_ep_train.py @@ -137,6 +137,39 @@ class EQBlock: return f +def anderson_relax(blk, z, xin, budget, eps, m=5, tol=1e-6): + """Anderson acceleration (type-II, window m) for the free-phase fixed point g(z)=z+eps*F(z). + PURE-SPEED opt-in (--fastfp): z* is solver-independent, but the T1-residual acquires 'solved' + semantics (resreg fires less) — use for probes/sweeps/ladder, not for reg-ablation arms. + Safeguard: if the AA candidate's residual is worse than the plain step's, take the plain step.""" + with torch.no_grad(): + sh = z.shape + X, G = [], [] # iterate / g(iterate) histories, flattened + for k in range(budget): + f = eps * blk.force(z, xin).detach() + g = z + f + r = f.reshape(-1) + if r.norm() < tol * (z.norm() + 1e-12): + return g.detach(), k + 1 + X.append(z.reshape(-1).clone()); G.append(g.reshape(-1).clone()) + if len(X) > m + 1: + X.pop(0); G.pop(0) + if len(X) >= 2: + dR = torch.stack([(G[i + 1] - X[i + 1]) - (G[i] - X[i]) for i in range(len(X) - 1)], 1) + try: + al = torch.linalg.lstsq(dR, r.unsqueeze(1)).solution.squeeze(1) + dG = torch.stack([G[i + 1] - G[i] for i in range(len(X) - 1)], 1) + cand = (G[-1].unsqueeze(1) - dG @ al.unsqueeze(1)).squeeze(1).reshape(sh) + rc = (eps * blk.force(cand, xin)).norm() + if rc < r.norm(): # safeguarded acceptance + z = cand + continue + except Exception: + pass + z = g + return z.detach(), budget + + def relax(blk, z, xin, steps, eps): cstep = getattr(blk, '_cstep', None) if cstep is not None and blk.fnoise == 0.0: # compiled pure-thick free-phase fast path @@ -146,6 +179,15 @@ def relax(blk, z, xin, steps, eps): return z.detach() blk._sdpa = getattr(blk, 'sdpa', False) # fused attention for the pure-forward loop only try: + if getattr(blk, 'fastfp', False) and steps >= 50 and blk.fnoise == 0.0: + zf, _ = anderson_relax(blk, z, xin, steps, eps) + return zf + k = getattr(blk, 'bf16polish', 0) # bf16 bulk + fp32 endpoint polish (last k steps) + if k and steps > k and blk.fnoise == 0.0: + with torch.no_grad(), torch.autocast('cuda', dtype=torch.bfloat16): + for _ in range(steps - k): + z = (z + eps * blk.force(z, xin)).float().detach() + steps = k for _ in range(steps): with torch.no_grad(): z = z + eps * blk.force(z, xin).detach() @@ -447,6 +489,8 @@ def main(): ap.add_argument('--holofast', action='store_true') # exact halved-jvp track (1.55x nudged phase; parity = FD noise floor) ap.add_argument('--sdpa', action='store_true') # fused flash attention in the no_grad relax loop ap.add_argument('--holoavg', action='store_true') # trend-aware stop + plateau-avg track (gate: 0.913->0.936 @t2sel160) + ap.add_argument('--fastfp', action='store_true') # Anderson-accelerated free phase (pure-speed opt-in; alters resreg semantics) + ap.add_argument('--bf16polish', type=int, default=0) # bf16 bulk relax + fp32 last-K polish (0=off; gate before use) ap.add_argument('--rt_final', type=float, default=0.0) # anneal res_target to this (0=off), 25%-75% of run ap.add_argument('--nudge_brake', type=float, default=0.0) # kappa: anchor spring during nudge (Tikhonov adjoint) ap.add_argument('--init_ckpt', type=str, default='') # warm-start weights from a saved ckpt @@ -524,6 +568,8 @@ def main(): blk.holofast = cfg.holofast blk.sdpa = cfg.sdpa blk.holoavg = cfg.holoavg + blk.fastfp = cfg.fastfp + blk.bf16polish = cfg.bf16polish blk.nbrake = cfg.nudge_brake blk.qknorm = cfg.qknorm if cfg.resinit != 1.0: # near-identity block at init (contractive) -> stable big-width start -- cgit v1.2.3