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

2 files changed, 11 insertions, 7 deletions

diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py
index db6d71a..a57e7b2 100644
--- a/ep_run/casc_eq_train.py
+++ b/ep_run/casc_eq_train.py
@@ -139,6 +139,8 @@ ap.add_argument('--relax_tol', type=float, default=0.0)   # >0: ADAPTIVE relax �
                                                           # signal), then one final graphed round. 0 = legacy fixed-K.
 ap.add_argument('--muon_mom', type=float, default=0.95)   # Muon momentum (late-SNR arm: 0.99 = ~10x noise averaging)
 ap.add_argument('--adam_b1', type=float, default=0.9)     # AdamW beta1 (late-SNR arm companion)
+ap.add_argument('--adam_b2', type=float, default=0.999)  # AdamW beta2 (LLM practice often 0.95)
+ap.add_argument('--muon_wd', type=float, default=0.0)    # decoupled wd on Muon matrices (sigma/drift-peak fix arm)
 ap.add_argument('--est', choices=['single', 'centered', 'richardson'], default='single')
                                                           # centered: [g(+b)+g(-b)]/2 (O(b^2) bias, 2x relax cost)
                                                           # richardson: 2g(b)-g(2b)   (O(b^2) bias, large-b friendly)
@@ -356,16 +358,17 @@ elif args.opt == 'muon':
                               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,
                               head_param=(W_out if args.untie or args.olmo2 else None),
-                              head_lr_mult=args.head_lr_mult)
+                              head_lr_mult=args.head_lr_mult, muon_wd=args.muon_wd)
 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()} | \
                   {id(p) for p in blocks.parameters() if p.ndim < 2} | {id(p) for p in ln_f.parameters()}
         opt = torch.optim.AdamW([
             {'params': [p for p in all_params if id(p) not in nodecay], 'weight_decay': args.wd},
-            {'params': [p for p in all_params if id(p) in nodecay], 'weight_decay': 0.0}], lr=args.lr)
+            {'params': [p for p in all_params if id(p) in nodecay], 'weight_decay': 0.0}], lr=args.lr,
+            betas=(args.adam_b1, args.adam_b2))
     else:
-        opt = torch.optim.AdamW(all_params, lr=args.lr, weight_decay=1e-4)
+        opt = torch.optim.AdamW(all_params, lr=args.lr, weight_decay=1e-4, betas=(args.adam_b1, args.adam_b2))
     if args.cosine:
         def _lrlam(s):
             if s < args.warmup: return (s + 1) / max(args.warmup, 1)
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},