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
|
"""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)
|
