RESULT 66: head被节流到60%(位移微探针:全族1.00-1.02唯W_out 0.60) — 用户等价LR假说命中并定位到单矩阵; Adam岛m/√v节流,对centered/fp32/K8全不敏感=EP头读共享结构; --head_lr_mult实装; 电池2(headmix诊断+hlr1.67/2.5补偿)在飞

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn

3 files changed, 37 insertions, 3 deletions

diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py
index 0cd5521..07e2ad4 100644
--- a/ep_run/casc_eq_train.py
+++ b/ep_run/casc_eq_train.py
@@ -49,6 +49,8 @@ ap.add_argument('--bsign_rand', action='store_true')  # random-sign beta per ste
 ap.add_argument('--bf16', action='store_true')           # cast model to bf16 (E-accumulation + tok_sigma stay fp32) — the x0.5 cost lever, GATE before production
 ap.add_argument('--amp', action='store_true')            # PROPER mixed precision: autocast(bf16) matmuls, fp32 params/states/d/E — amp_gate.py PASSED 2026-07-12 (cos 0.9682 vs fp32 0.9687); --bf16 naive-cast stays DEAD (state quantization, RESULT 11)
 ap.add_argument('--dtop_every', type=int, default=1)    # 1 = exact (DEFAULT, BP-parity); 2 = fast mode (~20% cheaper, ~4% CE tax at high lr)
+ap.add_argument('--head_lr_mult', type=float, default=1.0)  # W_out Adam-group LR multiplier (C768
+                                                          # head-throttle compensation, RESULT 66)
 ap.add_argument('--res_gate', type=float, default=0.02)  # legality residual threshold; 0.02 was calibrated
                                                           # at C512 — C768 ran half a schedule semi-converged
                                                           # UNDER it (R59b). Per-width rule: ~100x the healthy
@@ -313,7 +315,9 @@ if args.opt == 'muon':
     from muon import build_hybrid
     opt, sched = build_hybrid(blocks, all_params, args.lr, args.muon_lr, args.warmup,
                               muon_mom=args.muon_mom, adam_b1=args.adam_b1,
-                              total_steps=(args.steps if args.cosine else 0), lr_min_ratio=args.lr_min_ratio)
+                              total_steps=(args.steps if args.cosine else 0), lr_min_ratio=args.lr_min_ratio,
+                              head_param=(W_out if args.untie or args.olmo2 else None),
+                              head_lr_mult=args.head_lr_mult)
 else:
     if args.wd >= 0:   # OLMo2-style grouped decay: linear weights + head decay; embeddings/norm-gains none
         nodecay = {id(p) for p in tok.parameters()} | {id(p) for p in pos.parameters()} | \
diff --git a/ep_run/muon.py b/ep_run/muon.py
index d65a231..0806243 100644
--- a/ep_run/muon.py
+++ b/ep_run/muon.py
@@ -57,7 +57,7 @@ class MultiSched:
 
 
 def build_hybrid(blocks, other_params, lr_adamw, lr_muon, warmup, total_steps=0, lr_min_ratio=0.1,
-                 muon_mom=0.95, adam_b1=0.9):
+                 muon_mom=0.95, adam_b1=0.9, head_param=None, head_lr_mult=1.0):
     """Muon(2D block matrices) + AdamW(everything else). Scheds: linear warmup, then cosine decay to
     lr_min_ratio*peak if total_steps>0 (long runs), else constant after warmup (legacy).
     muon_mom/adam_b1: momentum knobs (late-SNR noise-averaging arms, 2026-07-13)."""
@@ -66,7 +66,13 @@ def build_hybrid(blocks, other_params, lr_adamw, lr_muon, warmup, total_steps=0,
     mat_ids = {id(p) for p in mats}
     rest = [p for p in other_params if id(p) not in mat_ids]
     om = Muon(mats, lr=lr_muon, momentum=muon_mom)
-    oa = torch.optim.AdamW(rest, lr=lr_adamw, weight_decay=1e-4, betas=(adam_b1, 0.999))
+    if head_param is not None and head_lr_mult != 1.0:
+        hid = id(head_param)
+        groups = [{'params': [p for p in rest if id(p) != hid], 'lr': lr_adamw},
+                  {'params': [p for p in rest if id(p) == hid], 'lr': lr_adamw * head_lr_mult}]
+        oa = torch.optim.AdamW(groups, lr=lr_adamw, weight_decay=1e-4, betas=(adam_b1, 0.999))
+    else:
+        oa = torch.optim.AdamW(rest, lr=lr_adamw, weight_decay=1e-4, betas=(adam_b1, 0.999))
     if total_steps > 0:
         def fn(s):
             if s < warmup: return (s + 1) / max(warmup, 1)
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()))