File size: 7,988 Bytes
fe397d3 | 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 | """Optimizers for Task 2. Rules: subclass torch.optim.Optimizer directly; no other torch.optim.* anywhere.
Paper: "Fantastic Pretraining Optimizers and Where to Find Them" (Wen et al., 2025), Appendix A has one algorithm per
optimizer. Every step() below follows the same pattern:
for each param group -> for each parameter p with a gradient:
g = p.grad (same shape as p), state = self.state[p] (dict that persists across steps, empty on step 1)
-> initialise state on the first step -> update the running statistics -> apply the update to p in place.
Hyperparameters come from `group` (the dict built from `defaults`, overridable per group). The training loop sets
group["lr"] every step (learning-rate schedule), so always read the lr from the group, never cache it.
"""
import torch
from torch.optim import Optimizer
def _adamw_update(p, g, state, group):
lr, (beta1, beta2), eps, wd = group["lr"], group["betas"], group["eps"], group["weight_decay"]
if len(state) == 0:
state["step"] = 0
state["exp_avg"] = torch.zeros_like(p)
state["exp_avg_sq"] = torch.zeros_like(p)
state["step"] += 1
t = state["step"]
exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
exp_avg.mul_(beta1).add_(g, alpha=1 - beta1) # m_t = beta_1 * m_t-1 + (1 - beta_1) * g_t
exp_avg_sq.mul_(beta2).addcmul_(g, g, value=1 - beta2) #
bc1 = 1 - beta1 ** t
bc2 = 1 - beta2 ** t
p.mul_(1 - lr * wd)
denom = (exp_avg_sq / bc2).sqrt_().add_(eps)
p.addcdiv_(exp_avg, denom, value=-lr / bc1)
class AdamW(Optimizer):
"""Category 1 (baseline). Appendix A, Algorithm 1."""
def __init__(self, params, lr=6e-4, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.1):
super().__init__(params, dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay))
@torch.no_grad()
def step(self, closure=None):
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()
for group in self.param_groups:
for p in group["params"]:
if p.grad is None:
continue
g, state = p.grad, self.state[p]
_adamw_update(p, g, state, group)
return loss
class NAdamW(Optimizer):
"""Category 2 (variance-reduced AdamW variants). Appendix A, Algorithm 2.
Same state as AdamW; the difference is how the momentum enters the update (Nesterov look-ahead)."""
def __init__(self, params, lr=6e-4, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.1):
super().__init__(params, dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay))
@torch.no_grad()
def step(self, closure=None):
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()
for group in self.param_groups:
lr, (beta1, beta2), eps, wd = group["lr"], group["betas"], group["eps"], group["weight_decay"]
for p in group["params"]:
if p.grad is None:
continue
g, state = p.grad, self.state[p]
if len(state) == 0:
state["step"] = 0
state["exp_avg"] = torch.zeros_like(p)
state["exp_avg_sq"] = torch.zeros_like(p)
state["step"] += 1
t = state["step"]
exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
exp_avg.mul_(beta1).add_(g, alpha=1 - beta1)
exp_avg_sq.mul_(beta2).addcmul_(g, g, value=1 - beta2)
bc1 = 1 - beta1 ** t
bc2 = 1 - beta2 ** t
nesterov_m = exp_avg.mul(beta1).add_(g, alpha=1 - beta1)
p.mul_(1 - lr * wd)
denom = (exp_avg_sq / bc2).sqrt_().add_(eps)
p.addcdiv_(nesterov_m, denom, value=-lr / bc1)
return loss
class Lion(Optimizer):
"""Category 3 (memory-efficient). Appendix A, Algorithm 3.
One state buffer per parameter; the update is a sign, so every coordinate moves by exactly lr
-> needs a much smaller lr (and usually larger weight decay) than AdamW."""
def __init__(self, params, lr=1e-4, betas=(0.9, 0.99), weight_decay=0.1):
super().__init__(params, dict(lr=lr, betas=betas, weight_decay=weight_decay))
@torch.no_grad()
def step(self, closure=None):
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()
for group in self.param_groups:
lr, (beta1, beta2), wd = group["lr"], group["betas"], group["weight_decay"]
for p in group["params"]:
if p.grad is None:
continue
g, state = p.grad, self.state[p]
if len(state) == 0:
state["exp_avg"] = torch.zeros_like(p)
exp_avg = state["exp_avg"]
update = exp_avg.mul(beta1).add_(g, alpha=1 - beta1).sign_()
p.mul_(1 - lr * wd)
p.add_(update, alpha=-lr)
exp_avg.mul_(beta2).add_(g, alpha=1 - beta2)
return loss
def zeropower_via_newtonschulz(G, steps):
"""G: (m, n) matrix -> matrix of the same shape with the same singular vectors and singular values pushed to ~1
(an approximate orthogonalisation), computed with `steps` Newton-Schulz iterations. Used by Muon."""
a, b, c = 3.4445, -4.7750, 2.0315
X = G.float()
X = X / (X.norm() + 1e-7)
transposed = G.size(0) > G.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
if transposed:
X = X.T
return X
class Muon(Optimizer):
"""Category 4 (matrix-based). Appendix A, Algorithm 8.
Two kinds of param groups (built in train.py):
use_muon=True -> hidden 2-D weight matrices (attention q/k/v/out, FFN linears): Muon update
use_muon=False -> embedding (= tied output layer), norm gains, biases: AdamW update (reuse your AdamW logic)
The two kinds use different learning rates (Muon's is much larger), set per group in train.py."""
def __init__(self, params, lr=0.02, momentum=0.95, nesterov=True, ns_steps=5, weight_decay=0.0,
betas=(0.9, 0.95), eps=1e-8):
super().__init__(params, dict(lr=lr, momentum=momentum, nesterov=nesterov, ns_steps=ns_steps,
weight_decay=weight_decay, betas=betas, eps=eps, use_muon=True))
@torch.no_grad()
def step(self, closure=None):
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()
for group in self.param_groups:
for p in group["params"]:
if p.grad is None:
continue
g, state = p.grad, self.state[p]
if group["use_muon"]:
lr, momentum, wd = group["lr"], group["momentum"], group["weight_decay"]
if len(state) == 0:
state["momentum_buffer"] = torch.zeros_like(p)
buf = state["momentum_buffer"]
buf.lerp_(g, 1 - momentum)
update = g.lerp(buf, momentum) if group["nesterov"] else buf
update = zeropower_via_newtonschulz(update, group["ns_steps"])
scale = max(1.0, p.size(0) / p.size(1)) ** 0.5
p.mul_(1 - lr * wd)
p.add_(update.to(p.dtype), alpha=-lr * scale)
else:
_adamw_update(p, g, state, group)
return loss
OPTIMIZERS = {"adamw": AdamW, "nadamw": NAdamW, "lion": Lion, "muon": Muon}
|