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/muon.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) (limited to 'ep_run/muon.py') 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