From 98d5ea09d1498f802f2ad06f5f67b385e228f063 Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Tue, 28 Jul 2026 13:20:34 -0500 Subject: =?UTF-8?q?RESULT=2066:=20head=E8=A2=AB=E8=8A=82=E6=B5=81=E5=88=B0?= =?UTF-8?q?60%(=E4=BD=8D=E7=A7=BB=E5=BE=AE=E6=8E=A2=E9=92=88:=E5=85=A8?= =?UTF-8?q?=E6=97=8F1.00-1.02=E5=94=AFW=5Fout=200.60)=20=E2=80=94=20?= =?UTF-8?q?=E7=94=A8=E6=88=B7=E7=AD=89=E4=BB=B7LR=E5=81=87=E8=AF=B4?= =?UTF-8?q?=E5=91=BD=E4=B8=AD=E5=B9=B6=E5=AE=9A=E4=BD=8D=E5=88=B0=E5=8D=95?= =?UTF-8?q?=E7=9F=A9=E9=98=B5;=20Adam=E5=B2=9Bm/=E2=88=9Av=E8=8A=82?= =?UTF-8?q?=E6=B5=81,=E5=AF=B9centered/fp32/K8=E5=85=A8=E4=B8=8D=E6=95=8F?= =?UTF-8?q?=E6=84=9F=3DEP=E5=A4=B4=E8=AF=BB=E5=85=B1=E4=BA=AB=E7=BB=93?= =?UTF-8?q?=E6=9E=84;=20--head=5Flr=5Fmult=E5=AE=9E=E8=A3=85;=20=E7=94=B5?= =?UTF-8?q?=E6=B1=A02(headmix=E8=AF=8A=E6=96=AD+hlr1.67/2.5=E8=A1=A5?= =?UTF-8?q?=E5=81=BF)=E5=9C=A8=E9=A3=9E?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn --- ep_run/muon.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) (limited to 'ep_run/muon.py') diff --git a/ep_run/muon.py b/ep_run/muon.py index d65a231..0806243 100644 --- a/ep_run/muon.py +++ b/ep_run/muon.py @@ -57,7 +57,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): + muon_mom=0.95, adam_b1=0.9, head_param=None, head_lr_mult=1.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).""" @@ -66,7 +66,13 @@ def build_hybrid(blocks, other_params, lr_adamw, lr_muon, warmup, total_steps=0, 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) - oa = torch.optim.AdamW(rest, lr=lr_adamw, weight_decay=1e-4, betas=(adam_b1, 0.999)) + 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}, + {'params': [p for p in rest if id(p) == hid], 'lr': lr_adamw * head_lr_mult}] + oa = torch.optim.AdamW(groups, lr=lr_adamw, weight_decay=1e-4, betas=(adam_b1, 0.999)) + else: + oa = torch.optim.AdamW(rest, lr=lr_adamw, weight_decay=1e-4, betas=(adam_b1, 0.999)) if total_steps > 0: def fn(s): if s < warmup: return (s + 1) / max(warmup, 1) -- cgit v1.2.3