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
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
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)
|
