"""Direct effective-step-size measurement (user hypothesis: high cos but shrunken equivalent
LR -> slow learning masquerading as plateau). Load two ckpts of the same tag 600 steps apart
(saved by the micro-runs), report per-family ||dW||/||W|| — the actual distance moved in
weight space. Compare EP vs BP micro-runs from the same start."""
import sys, torch
a, b, label = sys.argv[1], sys.argv[2], sys.argv[3]
ca, cb = (torch.load(p, map_location='cpu', weights_only=False) for p in (a, b))
fams = {'tok': None, 'wout': None}
out = {}
ta, tb = ca['tok']['weight'].float(), cb['tok']['weight'].float()
out['tok'] = float((tb-ta).norm()/ta.norm())
wa, wb = ca['wout'].float(), cb['wout'].float()
out['wout'] = float((wb-wa).norm()/wa.norm())
Ba, Bb = ca['blocks'], cb['blocks']
import re
groups = {'attn_bot': [], 'attn_top': [], 'ffn_bot': [], 'ffn_top': []}
for k in Ba:
    m = re.match(r'(\d+)\.(attn|ff)\.', k)
    if not m or Ba[k].dim() < 2: continue
    l = int(m.group(1)); fam = ('attn' if m.group(2)=='attn' else 'ffn') + ('_bot' if l < 6 else '_top')
    d = float((Bb[k].float()-Ba[k].float()).norm()/Ba[k].float().norm())
    groups[fam].append(d)
for k, v in groups.items(): out[k] = sum(v)/len(v)
print(label + ': ' + ' '.join(f'{k} {v:.5f}' for k, v in out.items()))