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