diff options
Diffstat (limited to 'ep_run/lt_ep_train.py')
| -rw-r--r-- | ep_run/lt_ep_train.py | 17 |
1 files changed, 16 insertions, 1 deletions
diff --git a/ep_run/lt_ep_train.py b/ep_run/lt_ep_train.py index 2807be3..76bb9bb 100644 --- a/ep_run/lt_ep_train.py +++ b/ep_run/lt_ep_train.py @@ -81,6 +81,20 @@ class EQBlock: def Emem(self, z): return -(F.relu(z @ self.Wm) ** 2).sum() + def tforce_sdpa(self, z, xin): # tforce with fused flash attention (compiled fast path only; + h1 = F.layer_norm(z, (self.C,), self.ln1g, self.ln1b) # grad paths keep the manual attn) + h2 = F.layer_norm(z, (self.C,), self.ln2g, self.ln2b) + B = z.size(0) + q = (h1 @ self.WQ).view(B, self.T, self.H, self.dh).transpose(1, 2) + k = (h1 @ self.WK).view(B, self.T, self.H, self.dh).transpose(1, 2) + v = (h1 @ self.WV).view(B, self.T, self.H, self.dh).transpose(1, 2) + if getattr(self, 'qknorm', False): + 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) + att = F.scaled_dot_product_attention(q, k, v, is_causal=True).transpose(1, 2).reshape(B, self.T, self.C) @ self.WO + ff = F.gelu(h2 @ self.fc + self.fcb, approximate='tanh') @ self.pj + self.pjb + return -(z - xin) + att + ff - self.c * z + def tforce(self, z, xin): # pure thick force (no grad machinery) -> torch.compile h1 = F.layer_norm(z, (self.C,), self.ln1g, self.ln1b) h2 = F.layer_norm(z, (self.C,), self.ln2g, self.ln2b) @@ -523,7 +537,8 @@ def main(): blk._cstep = None if cfg.compile and cfg.attn_mode == 'thick': _ee = cfg.eps - blk._cstep = torch.compile(lambda z, xin: z + _ee * blk.tforce(z, xin)) + _tf = blk.tforce_sdpa if cfg.sdpa else blk.tforce # sdpa baked into the compiled graph, NO ambient flags + blk._cstep = torch.compile(lambda z, xin: z + _ee * _tf(z, xin)) mis = None if cfg.wmis > 0: # fixed fabrication mismatch (same devices all run) gm = torch.Generator().manual_seed(1234) |