wd臂(用户批准, 预注册后发射): Muon矩阵组加解耦衰减--muon_wd, C512全程臂链在bp_s3后HEADmaster

预测冻结在启动前(见wd_arm.sh注释): P1 sig_late降>15%(基线521), P2 drift峰<0.015, P3 尾窗
3.4003±0.010内或更好。P1&P2中+P3不劣 => wd入270M+冻结配方; P3劣 => 降档或改Hyperball控σ。
背景: drift标度分析(08-07)显示末段均值近平坦(C^0.10)但峰值顶端加速(C768=0.041=C512的2.7×),
根因σ∝C增长, 外推C1536-2048触0.5守卫线; wd是已知洞(muon分支无衰减)与该风险的连接点。
附: --adam_b2实装+Adam电池5臂在农场(β2=0.95从未测过)。

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

1 files changed, 5 insertions, 4 deletions

diff --git a/ep_run/muon.py b/ep_run/muon.py
index 86c81d3..0633b2e 100644
--- a/ep_run/muon.py
+++ b/ep_run/muon.py
@@ -17,8 +17,8 @@ def newton_schulz(G, steps=5, eps=1e-7):
 
 
 class Muon(torch.optim.Optimizer):
-    def __init__(self, params, lr=0.02, momentum=0.95, ns_steps=5, nesterov=True):
-        super().__init__(params, dict(lr=lr, momentum=momentum, ns_steps=ns_steps, nesterov=nesterov))
+    def __init__(self, params, lr=0.02, momentum=0.95, ns_steps=5, nesterov=True, wd=0.0):
+        super().__init__(params, dict(lr=lr, momentum=momentum, ns_steps=ns_steps, nesterov=nesterov, wd=wd))
 
     @torch.no_grad()
     def step(self, closure=None):
@@ -34,6 +34,7 @@ class Muon(torch.optim.Optimizer):
                 if u.ndim == 2:
                     u = newton_schulz(u, group['ns_steps'])
                     u = u * max(1.0, u.size(0) / u.size(1)) ** 0.5   # rms-matched scaling
+                if group['wd'] > 0: p.mul_(1 - group['lr'] * group['wd'])   # decoupled decay (Moonshot: required for scale)
                 p.add_(u, alpha=-group['lr'])
 
 
@@ -57,7 +58,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, head_param=None, head_lr_mult=1.0):
+                 muon_mom=0.95, adam_b1=0.9, head_param=None, head_lr_mult=1.0, muon_wd=0.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)."""
@@ -65,7 +66,7 @@ def build_hybrid(blocks, other_params, lr_adamw, lr_muon, warmup, total_steps=0,
     mats = [p for p in blocks.parameters() if p.ndim == 2]
     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)
+    om = Muon(mats, lr=lr_muon, momentum=muon_mom, wd=muon_wd)
     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},