diff options
Diffstat (limited to 'ep_run/muon.py')
| -rw-r--r-- | ep_run/muon.py | 9 |
1 files changed, 5 insertions, 4 deletions
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}, |