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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
|
"""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/state_dict 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)
def state_dict(self):
return [o.state_dict() for o in self.optimizers]
def load_state_dict(self, sds):
for o, sd in zip(self.optimizers, sds): o.load_state_dict(sd)
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, 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(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)."""
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, momentum=muon_mom)
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)
p = min(1.0, (s - warmup) / max(1, total_steps - warmup))
return lr_min_ratio + 0.5 * (1 - lr_min_ratio) * (1 + _m.cos(_m.pi * p))
else:
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)
class Lion(torch.optim.Optimizer):
"""Lion (Chen et al. 2023): sign of the beta1-mixed momentum; decoupled wd."""
def __init__(self, params, lr=3e-4, betas=(0.9, 0.99), wd=0.0):
super().__init__(params, dict(lr=lr, betas=betas, wd=wd))
@torch.no_grad()
def step(self, closure=None):
for g_ in self.param_groups:
b1, b2 = g_['betas']
for p in g_['params']:
if p.grad is None: continue
st = self.state[p]
if 'm' not in st: st['m'] = torch.zeros_like(p)
m = st['m']
u = (b1 * m + (1 - b1) * p.grad).sign_()
if g_['wd'] > 0: p.mul_(1 - g_['lr'] * g_['wd'])
p.add_(u, alpha=-g_['lr'])
m.mul_(b2).add_(p.grad, alpha=1 - b2)
class OLion(torch.optim.Optimizer):
"""OLion (arXiv:2602.01105): Lion-style momentum -> Newton-Schulz -> entrywise sign,
RMS alignment gamma (||sign||_F = sqrt(numel) exactly, so D = gamma * sign(Q)), decoupled wd.
2D params only; caller routes 1D elsewhere. betas default to Lion's (0.9, 0.99) — the paper's
defaults were not in the pages we read; flagged as an assumption in the battery notes."""
def __init__(self, params, lr=1e-3, betas=(0.9, 0.99), gamma=0.2, ns_steps=5, wd=0.0):
super().__init__(params, dict(lr=lr, betas=betas, gamma=gamma, ns_steps=ns_steps, wd=wd))
@torch.no_grad()
def step(self, closure=None):
for g_ in self.param_groups:
b1, b2 = g_['betas']
for p in g_['params']:
if p.grad is None: continue
st = self.state[p]
if 'm' not in st: st['m'] = torch.zeros_like(p)
m = st['m']
m.mul_(b2).add_(p.grad, alpha=1 - b2) # slow momentum
gt = (1 - b1) * p.grad + b1 * m # Nesterov mix
q = newton_schulz(gt, g_['ns_steps'])
if g_['wd'] > 0: p.mul_(1 - g_['lr'] * g_['wd'])
p.add_(q.sign_(), alpha=-g_['lr'] * g_['gamma'])
class Adafactor2D(torch.optim.Optimizer):
"""Minimal Shazeer-Stern factored Adam for 2D params: row/col second-moment statistics,
RMS-1 update clipping, no relative-step magic (external lr + schedule). The analog-native
Adam per the 07-11 BoM audit (row/col stats = AGC channels)."""
def __init__(self, params, lr=1e-3, beta2=0.999, eps=1e-30, clip=1.0, wd=0.0):
super().__init__(params, dict(lr=lr, beta2=beta2, eps=eps, clip=clip, wd=wd))
@torch.no_grad()
def step(self, closure=None):
for g_ in self.param_groups:
for p in g_['params']:
if p.grad is None: continue
g = p.grad
st = self.state[p]
if 'r' not in st:
st['r'] = torch.zeros(g.shape[0], device=g.device)
st['c'] = torch.zeros(g.shape[1], device=g.device)
r, c = st['r'], st['c']
g2 = g.float().pow(2) + g_['eps']
r.mul_(g_['beta2']).add_(g2.mean(1), alpha=1 - g_['beta2'])
c.mul_(g_['beta2']).add_(g2.mean(0), alpha=1 - g_['beta2'])
v = r[:, None] * c[None, :] / max(float(r.mean()), g_['eps'])
u = g / v.sqrt().to(g.dtype)
rms = float(u.pow(2).mean().sqrt())
if rms > g_['clip']: u = u * (g_['clip'] / rms)
if g_['wd'] > 0: p.mul_(1 - g_['lr'] * g_['wd'])
p.add_(u, alpha=-g_['lr'])
def build_alt(opt_name, blocks, other_params, lr, warmup, total_steps=0, lr_min_ratio=0.1,
lr_matrix=None, wd=0.0):
"""Screening-tier builder for the optimizer price list: OPT on block matrices + AdamW on the
rest (same split as build_hybrid so arms differ only in the matrix rule). sgdm applies SGD to
everything (no split) — the fully-local baseline."""
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]
lm = lr_matrix if lr_matrix is not None else lr
if opt_name == 'sgdm':
om = torch.optim.SGD(mats + rest, lr=lm, momentum=0.95, nesterov=True)
opts = [om]
else:
om = {'lion': lambda: Lion(mats, lr=lm, wd=wd),
'olion': lambda: OLion(mats, lr=lm, wd=wd),
'adafactor': lambda: Adafactor2D(mats, lr=lm, wd=wd)}[opt_name]()
oa = torch.optim.AdamW(rest, lr=lr, weight_decay=1e-4)
opts = [om, oa]
if total_steps > 0:
def fn(s):
if s < warmup: return (s + 1) / max(warmup, 1)
pr = min(1.0, (s - warmup) / max(1, total_steps - warmup))
return lr_min_ratio + 0.5 * (1 - lr_min_ratio) * (1 + _m.cos(_m.pi * pr))
else:
fn = lambda s: min(1.0, (s + 1) / max(warmup, 1))
scheds = [torch.optim.lr_scheduler.LambdaLR(o, fn) for o in opts]
return MultiOpt(opts), MultiSched(scheds)
|
