From 75ff326dcb40cd960abd56e9c9c18a45d9e5e2c2 Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Sun, 5 Jul 2026 05:09:36 -0500 Subject: =?UTF-8?q?--sdpa:=20fused=20flash=20attention=20in=20the=20no=5Fg?= =?UTF-8?q?rad=20relax=20loop=20=E2=80=94=201.45x=20free=20phase,=20z*=20p?= =?UTF-8?q?arity=204e-7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Scoped via blk._sdpa set only inside relax()'s loop (grad paths jvp/vjp/resreg keep the manual attention: no forward-mode-through-flash risk). Combined with --holofast: ~1.51x full-step exact-math tier. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn --- ep_run/lt_ep_train.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) (limited to 'ep_run/lt_ep_train.py') diff --git a/ep_run/lt_ep_train.py b/ep_run/lt_ep_train.py index e7155d3..1307702 100644 --- a/ep_run/lt_ep_train.py +++ b/ep_run/lt_ep_train.py @@ -63,6 +63,9 @@ class EQBlock: if getattr(self, 'qknorm', False): # Qwen3-style q/k RMSNorm: bounds logits, tames J q = q * torch.rsqrt(q.pow(2).mean(-1, keepdim=True) + 1e-6) k = k * torch.rsqrt(k.pow(2).mean(-1, keepdim=True) + 1e-6) + if getattr(self, '_sdpa', False): # fused flash path — no_grad relax/eval only (same + o = F.scaled_dot_product_attention(q, k, v, is_causal=True) # scale 1/sqrt(dh), same causal mask) + return o.transpose(1, 2).reshape(B, self.T, self.C) @ self.WO a = (q @ k.transpose(-2, -1)) / math.sqrt(self.dh) a = torch.softmax(a.masked_fill(~self.cmask, float('-inf')), -1) return (a @ v).transpose(1, 2).reshape(B, self.T, self.C) @ self.WO @@ -127,9 +130,13 @@ def relax(blk, z, xin, steps, eps): for _ in range(steps): z = cstep(z, xin) return z.detach() - for _ in range(steps): - with torch.no_grad(): - z = z + eps * blk.force(z, xin).detach() + blk._sdpa = getattr(blk, 'sdpa', False) # fused attention for the pure-forward loop only + try: + for _ in range(steps): + with torch.no_grad(): + z = z + eps * blk.force(z, xin).detach() + finally: + blk._sdpa = False # grad paths (jvp/vjp/resreg graphs) stay on manual attn return z.detach() @@ -421,6 +428,7 @@ def main(): ap.add_argument('--navg', type=int, default=1) # restart-averaged contrast estimates per update ap.add_argument('--track', action='store_true') # common-mode-tracking AEP correction 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('--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 @@ -496,6 +504,7 @@ def main(): blk.navg = cfg.navg blk.track = cfg.track blk.holofast = cfg.holofast + blk.sdpa = cfg.sdpa 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