diff options
Diffstat (limited to 'ep_run/holo_fast_gate.py')
| -rw-r--r-- | ep_run/holo_fast_gate.py | 25 |
1 files changed, 25 insertions, 0 deletions
diff --git a/ep_run/holo_fast_gate.py b/ep_run/holo_fast_gate.py new file mode 100644 index 0000000..2a5f39b --- /dev/null +++ b/ep_run/holo_fast_gate.py @@ -0,0 +1,25 @@ +"""Ship-gate for --holofast: gradient-level equivalence. cos(EP,BPTT) with orig vs fast track on the +same batches (s2000, track path, holo=2 t2sel=40). Ship iff the two cos columns are statistically +indistinguishable (the a-level 45% parity gap is FD noise; what matters is the gradient direction).""" +import torch +import lt_ep_train as L +from diag_cos import cos_ep_bptt + +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)) +blk.track = True + +print(f"{'batch':>5} {'cos_orig':>9} {'cos_fast':>9} {'res':>9}", flush=True) +torch.manual_seed(11) +batches = [L.get_batch('train', 24, 256) for _ in range(3)] +for b, (idx, y) in enumerate(batches): + blk.holofast = False + c0, r0 = cos_ep_bptt(blk, idx, y, 150, 20, 0.1, 0.02, holo=2, hr=0.02, t2sel=40) + blk.holofast = True + c1, r1 = cos_ep_bptt(blk, idx, y, 150, 20, 0.1, 0.02, holo=2, hr=0.02, t2sel=40) + print(f"{b:>5} {c0:>9.4f} {c1:>9.4f} {r0:>9.2e}", flush=True) +print("HOLO_FAST_GATE_DONE", flush=True) |