diff options
Diffstat (limited to 'ep_run')
| -rw-r--r-- | ep_run/lt_ep_train.py | 15 | ||||
| -rw-r--r-- | ep_run/sdpa_gate.log | 4 | ||||
| -rw-r--r-- | ep_run/sdpa_gate.py | 32 |
3 files changed, 48 insertions, 3 deletions
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 diff --git a/ep_run/sdpa_gate.log b/ep_run/sdpa_gate.log new file mode 100644 index 0000000..88bb685 --- /dev/null +++ b/ep_run/sdpa_gate.log @@ -0,0 +1,4 @@ +manual: res=9.981e+00 val=3.1589 relax150=1.650s + sdpa: res=9.981e+00 val=3.1056 relax150=1.140s +z* rel-diff=3.99e-07 speed=1.45x +SDPA_GATE_DONE diff --git a/ep_run/sdpa_gate.py b/ep_run/sdpa_gate.py new file mode 100644 index 0000000..0ea4f60 --- /dev/null +++ b/ep_run/sdpa_gate.py @@ -0,0 +1,32 @@ +"""Ship-gate for --sdpa (fused flash attention in the no_grad relax loop): z* parity + res + val + timing. +Grad paths untouched by construction (the _sdpa flag is scoped to relax's loop), so no BPTT gate needed.""" +import time, torch +import lt_ep_train as L + +torch.manual_seed(0) +blk = L.EQBlock(512, 16, 256, 256, c=1.0, attn_mode='thick'); blk.qknorm = True +ck = torch.load('runs/redx_traj/s2000.pt', 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() + +out = {} +for name in ('manual', 'sdpa'): + blk.sdpa = (name == 'sdpa') + z = L.relax(blk, xin.clone(), xin, 150, 0.1) # warmup + result + res = (L.relax(blk, z, xin, 1, 0.1) - z).norm().item() + val = L.evaluate(blk, 150, 0.1, nb=4) + ts = [] + for _ in range(3): + torch.cuda.synchronize(); t = time.time() + L.relax(blk, xin.clone(), xin, 150, 0.1) + torch.cuda.synchronize(); ts.append(time.time() - t) + out[name] = (z, res, val, min(ts)) + print(f"{name:>6}: res={res:.3e} val={val:.4f} relax150={min(ts):.3f}s", flush=True) + +zd = ((out['sdpa'][0] - out['manual'][0]).norm() / (out['manual'][0].norm() + 1e-12)).item() +print(f"z* rel-diff={zd:.2e} speed={out['manual'][3]/out['sdpa'][3]:.2f}x", flush=True) +print("SDPA_GATE_DONE", flush=True) |