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

1 files changed, 8 insertions, 2 deletions

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)