1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
|
"""Muon optimizer (Newton-Schulz orthogonalized momentum) + hybrid helpers.
Convention: Muon on 2D hidden matrices, AdamW on everything else (emb/pos/LN/bias)."""
import torch
def newton_schulz(G, steps=5, eps=1e-7):
"""approximate polar factor of G via the quintic NS iteration (Keller Jordan coefficients)."""
a, b, c = 3.4445, -4.7750, 2.0315
X = G / (G.norm() + eps)
transposed = X.size(0) > X.size(1)
if transposed: X = X.T
for _ in range(steps):
A = X @ X.T
B = b * A + c * (A @ A)
X = a * X + B @ X
return X.T if transposed else X
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))
@torch.no_grad()
def step(self, closure=None):
for group in self.param_groups:
for p in group['params']:
if p.grad is None: continue
g = p.grad
st = self.state[p]
if 'mom' not in st: st['mom'] = torch.zeros_like(g)
buf = st['mom']
buf.mul_(group['momentum']).add_(g)
u = g.add(buf, alpha=group['momentum']) if group['nesterov'] else buf
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
p.add_(u, alpha=-group['lr'])
class MultiOpt:
"""duck-typed bundle of optimizers (step/zero_grad API-compatible)."""
def __init__(self, opts): self.optimizers = opts
def step(self):
for o in self.optimizers: o.step()
def zero_grad(self, set_to_none=True):
for o in self.optimizers: o.zero_grad(set_to_none=set_to_none)
class MultiSched:
def __init__(self, scheds): self.scheds = scheds
def step(self):
for s in self.scheds: s.step()
def build_hybrid(blocks, other_params, lr_adamw, lr_muon, warmup):
"""Muon(2D block matrices) + AdamW(everything else), with linear-warmup scheds for both."""
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)
fn = lambda s: min(1.0, (s + 1) / max(warmup, 1))
scheds = [torch.optim.lr_scheduler.LambdaLR(om, fn), torch.optim.lr_scheduler.LambdaLR(oa, fn)]
return MultiOpt([om, oa]), MultiSched(scheds)
|
