diff options
Diffstat (limited to 'ep_run/probe_stepdisp.py')
| -rw-r--r-- | ep_run/probe_stepdisp.py | 24 |
1 files changed, 24 insertions, 0 deletions
diff --git a/ep_run/probe_stepdisp.py b/ep_run/probe_stepdisp.py new file mode 100644 index 0000000..2d86867 --- /dev/null +++ b/ep_run/probe_stepdisp.py @@ -0,0 +1,24 @@ +"""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())) |
