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