From 643983aec870af4ba7e71eafc6c1ba37b2348730 Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Mon, 13 Jul 2026 23:27:01 -0500 Subject: Estimator arms machinery: --est centered/richardson (two-pass, O(b^2) bias), --muon_mom/--adam_b1 knobs; smoke passed (resume+centered cos 0.998) Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn --- ep_run/muon.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) (limited to 'ep_run/muon.py') diff --git a/ep_run/muon.py b/ep_run/muon.py index 7281811..d65a231 100644 --- a/ep_run/muon.py +++ b/ep_run/muon.py @@ -56,15 +56,17 @@ class MultiSched: for s in self.scheds: s.step() -def build_hybrid(blocks, other_params, lr_adamw, lr_muon, warmup, total_steps=0, lr_min_ratio=0.1): +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(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).""" + 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).""" import math as _m 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) - oa = torch.optim.AdamW(rest, lr=lr_adamw, weight_decay=1e-4) + 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 total_steps > 0: def fn(s): if s < warmup: return (s + 1) / max(warmup, 1) -- cgit v1.2.3