1 files changed, 15 insertions, 0 deletions
diff --git a/ep_run/spec_check.py b/ep_run/spec_check.py
new file mode 100644
index 0000000..be633d5
--- /dev/null
+++ b/ep_run/spec_check.py
@@ -0,0 +1,15 @@
+import torch, pickle
+from pathlib import Path
+import lt_ep_train as L
+L.DD = Path('data/tinystories_bpe'); L.vocab = pickle.load(open(L.DD/'meta.pkl','rb'))['vocab_size']
+blk = L.EQBlock(512,16,256,256, s=1.0, c=1.0, attn_mode='thick'); blk.qknorm=True
+ck = torch.load('runs/bptt_spec_snapshot.pt', map_location='cpu')
+with torch.no_grad():
+    for p,w in zip(blk.allp, ck['allp']): p.copy_(w.to('cpu').float())
+print(f"# BPTT ckpt step={ck.get('step')} best_val={ck.get('best')}")
+print(f"# specnorm cap = 0.9")
+for name in ['WQ','WK','WV','WO','fc','pj']:
+    W = getattr(blk, name).float()
+    s = torch.linalg.svdvals(W)[0].item()
+    flag = "  <-- EXCEEDS cap" if s > 0.9 else ""
+    print(f"{name}{tuple(W.shape)}: sigma_max={s:.3f}{flag}")