Tier-3 gates: Anderson math-yes/impl-no (parked for v2); bf16polish UNSAFE near-edge (eval-only); dp_ep.py ready

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

4 files changed, 164 insertions, 0 deletions

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