1 files changed, 64 insertions, 0 deletions
diff --git a/ep_run/probe_specaudit.py b/ep_run/probe_specaudit.py
new file mode 100644
index 0000000..041d068
--- /dev/null
+++ b/ep_run/probe_specaudit.py
@@ -0,0 +1,64 @@
+"""Weight-space spectral audit: which loop-gain component separates the plain vs cent
+lineages before the 196k ceiling collapse? Per-block sigma of Wq/Wk/Wv/proj/w1/w3/w2 and
+rms of qk-norm & post-norm gains, at matched steps. Pure CPU, reads the 5k ckpt ladders."""
+import torch, sys
+
+STEPS_BOTH = [150000, 170000, 185000, 190000, 195000, 200000]
+STEPS_PLAIN_EXTRA = [210000, 220000, 230000]
+L = 12
+
+def sig(w):
+    return float(torch.linalg.svdvals(w.float())[0])
+
+def rms(g):
+    return float(g.float().pow(2).mean().sqrt())
+
+def audit(path):
+    ck = torch.load(path, map_location='cpu', weights_only=False)
+    b = ck['blocks']
+    rows = []
+    for l in range(L):
+        qkv = b[f'{l}.attn.qkv.weight'].float()
+        wq, wk, wv = qkv[0:512], qkv[512:1024], qkv[1024:1536]
+        rows.append(dict(
+            q=sig(wq), k=sig(wk), v=sig(wv), pr=sig(b[f'{l}.attn.proj.weight']),
+            w1=sig(b[f'{l}.ff.w1.weight']), w3=sig(b[f'{l}.ff.w3.weight']),
+            w2=sig(b[f'{l}.ff.w2.weight']),
+            qg=rms(b[f'{l}.attn.qn.g']), kg=rms(b[f'{l}.attn.kn.g']),
+            na=rms(b[f'{l}.na.g']), nf=rms(b[f'{l}.nf.g'])))
+    wout = sig(ck['wout'])
+    return rows, wout
+
+def summarize(tag, step, rows, wout):
+    # aggregates: mean over blocks, top-half mean (6-11), max block, plus the derived products
+    def agg(key, blks):
+        vals = [rows[l][key] for l in blks]
+        return sum(vals) / len(vals)
+    bot, top = range(0, 6), range(6, 12)
+    logit = [rows[l]['qg'] * rows[l]['kg'] for l in range(L)]           # qk-norm logit scale
+    vpath = [rows[l]['v'] * rows[l]['pr'] for l in range(L)]            # attn value-path gain
+    fpath = [max(rows[l]['w1'], rows[l]['w3']) * rows[l]['w2'] for l in range(L)]
+    print(f'{tag} s{step//1000:>3}k | logit bot {agg("qg",bot)*agg("kg",bot):6.3f} top {agg("qg",top)*agg("kg",top):6.3f} max {max(logit):6.3f}(b{logit.index(max(logit))}) '
+          f'| vpath top {sum(vpath[6:])/6:7.2f} max {max(vpath):7.2f}(b{vpath.index(max(vpath))}) '
+          f'| ffn top {sum(fpath[6:])/6:7.2f} max {max(fpath):7.2f}(b{fpath.index(max(fpath))}) '
+          f'| na top {agg("na",top):5.3f} nf top {agg("nf",top):5.3f} | wout {wout:6.1f}', flush=True)
+    return dict(logit=logit, vpath=vpath, fpath=fpath)
+
+res = {}
+for lineage, steps in [('plain', STEPS_BOTH + STEPS_PLAIN_EXTRA), ('cent', STEPS_BOTH + STEPS_PLAIN_EXTRA)]:
+    for s in steps:
+        p = f'runs/fw72m_{lineage}_s{s}.pt'
+        try:
+            rows, wout = audit(p)
+        except FileNotFoundError:
+            print(f'{lineage} s{s}: MISSING', flush=True); continue
+        res[(lineage, s)] = summarize(lineage, s, rows, wout)
+
+# per-block divergence table at 195k (last matched healthy step)
+if ('plain', 195000) in res and ('cent', 195000) in res:
+    print('\nper-block plain/cent ratio at 195k (the pre-collapse fingerprint):', flush=True)
+    p, c = res[('plain', 195000)], res[('cent', 195000)]
+    print('blk  logit-ratio  vpath-ratio  ffn-ratio')
+    for l in range(L):
+        print(f'{l:3d}  {p["logit"][l]/c["logit"][l]:10.3f}  {p["vpath"][l]/c["vpath"][l]:10.3f}  {p["fpath"][l]/c["fpath"][l]:9.3f}')
+print('AUDIT_DONE', flush=True)