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, 6 insertions, 3 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)