1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
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)
|
