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
33
34
35
36
37
38
39
40
|
"""Gate for --fastfp (Anderson z*) and --bf16polish: endpoint parity + residual + eval-count/time vs the
150-step Euler reference, on two operators (deeply-trained fast-adaptive + near-edge s2000 — the hard
case: Anderson historically fails on cycling ops, s2000's marginal band is the stress test).
GPU shared with abl_pair — timings are relative, parity/evals exact."""
import time, torch
import lt_ep_train as L
for name, path in (('s2000', 'runs/redx_traj/s2000.pt'), ('fast-adaptive', 'runs/ep_fast_adaptive.pt')):
torch.manual_seed(0)
blk = L.EQBlock(512, 16, 256, 256, c=1.0, attn_mode='thick'); blk.qknorm = True
ck = torch.load(path, 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()
torch.cuda.synchronize(); t = time.time()
z_ref = L.relax(blk, xin.clone(), xin, 150, 0.1)
torch.cuda.synchronize(); t_ref = time.time() - t
r_ref = (L.relax(blk, z_ref, xin, 1, 0.1) - z_ref).norm().item()
print(f"[{name}] euler150 : {t_ref:5.2f}s res={r_ref:.2e}", flush=True)
torch.cuda.synchronize(); t = time.time()
z_aa, evals = L.anderson_relax(blk, xin.clone(), xin, 150, 0.1)
torch.cuda.synchronize(); t_aa = time.time() - t
r_aa = (L.relax(blk, z_aa, xin, 1, 0.1) - z_aa).norm().item()
zd = ((z_aa - z_ref).norm() / (z_ref.norm() + 1e-12)).item()
print(f"[{name}] anderson : {t_aa:5.2f}s res={r_aa:.2e} evals={evals} z-diff={zd:.2e}", flush=True)
blk.bf16polish = 20
torch.cuda.synchronize(); t = time.time()
z_bf = L.relax(blk, xin.clone(), xin, 150, 0.1)
torch.cuda.synchronize(); t_bf = time.time() - t
blk.bf16polish = 0
r_bf = (L.relax(blk, z_bf, xin, 1, 0.1) - z_bf).norm().item()
zdb = ((z_bf - z_ref).norm() / (z_ref.norm() + 1e-12)).item()
print(f"[{name}] bf16+20 : {t_bf:5.2f}s res={r_bf:.2e} z-diff={zdb:.2e}", flush=True)
print("FASTFP_GATE_DONE", flush=True)
|
