From bd66ffc989eca87ebd2ec7350bdd0961408bec6b Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Wed, 5 Aug 2026 19:23:25 -0500 Subject: =?UTF-8?q?wd=E8=87=82(=E7=94=A8=E6=88=B7=E6=89=B9=E5=87=86,=20?= =?UTF-8?q?=E9=A2=84=E6=B3=A8=E5=86=8C=E5=90=8E=E5=8F=91=E5=B0=84):=20Muon?= =?UTF-8?q?=E7=9F=A9=E9=98=B5=E7=BB=84=E5=8A=A0=E8=A7=A3=E8=80=A6=E8=A1=B0?= =?UTF-8?q?=E5=87=8F--muon=5Fwd,=20C512=E5=85=A8=E7=A8=8B=E8=87=82?= =?UTF-8?q?=E9=93=BE=E5=9C=A8bp=5Fs3=E5=90=8E?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 预测冻结在启动前(见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 Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn --- ep_run/casc_eq_train.py | 9 ++++++--- ep_run/muon.py | 9 +++++---- 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}, -- cgit v1.2.3