diff options
Diffstat (limited to 'ep_run/lt_ep_train.py')
| -rw-r--r-- | ep_run/lt_ep_train.py | 46 |
1 files changed, 46 insertions, 0 deletions
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 |