File size: 6,959 Bytes
5407bdc | 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 | """Correctness checks for the from-scratch optimizers.
These are the checks the report cites. Run them with:
python -m src.part2.verify_optimizers
The most important one is the first: AdamW is the only optimizer here with an
authoritative reference implementation to compare against, so if this file's
AdamW matches ``torch.optim.AdamW`` bit-for-bit, the shared scaffolding
(parameter groups, state handling, decoupled decay, bias correction) is known
good and the other four differ from it only in their documented ways.
"""
from __future__ import annotations
import torch
import torch.nn.functional as F
from .optimizers import OPTIMIZERS, Lion, MARS, SophiaG
def _model(seed=0):
torch.manual_seed(seed)
return torch.nn.Sequential(torch.nn.Linear(16, 32), torch.nn.GELU(),
torch.nn.Linear(32, 4))
def check_adamw_matches_torch(steps=25, tol=1e-6):
"""Ours vs torch.optim.AdamW from identical init on identical batches."""
a, b = _model(), _model()
oa = OPTIMIZERS["adamw"](a.parameters(), lr=1e-3, betas=(0.9, 0.95),
weight_decay=0.1)
ob = torch.optim.AdamW(b.parameters(), lr=1e-3, betas=(0.9, 0.95),
weight_decay=0.1, eps=1e-8)
torch.manual_seed(1)
for x in [torch.randn(8, 16) for _ in range(steps)]:
for m, o in ((a, oa), (b, ob)):
o.zero_grad()
m(x).pow(2).mean().backward()
o.step()
d = max((p - q).abs().max().item()
for p, q in zip(a.parameters(), b.parameters()))
ok = d < tol
print(f"[{'PASS' if ok else 'FAIL'}] AdamW == torch.optim.AdamW after "
f"{steps} steps: max|diff| = {d:.2e}")
return ok
def check_all_reduce_loss(steps=60):
"""Every optimizer must actually optimise, and report its state memory."""
ok = True
print(f"\n{'optimizer':10s} {'loss[0]':>9s} {'loss[-1]':>9s} "
f"{'state/param':>12s}")
lrs = {"sophia": 3e-3, "muon": 2e-2, "lion": 3e-4}
n_param = sum(p.numel() for p in _model().parameters())
for name, cls in OPTIMIZERS.items():
m = _model()
o = cls(m.parameters(), **({"lr": lrs[name]} if name in lrs else {}))
torch.manual_seed(1)
first = None
for i, x in enumerate([torch.randn(8, 16) for _ in range(steps)]):
o.zero_grad()
loss = m(x).pow(2).mean()
loss.backward()
o.step(**({"bs_scale": 8.0} if isinstance(o, SophiaG) else {}))
if i == 0:
first = loss.item()
mult = o.memory_bytes() / (n_param * 4)
good = loss.item() < first
ok &= good
print(f"{name:10s} {first:9.4f} {loss.item():9.4f} {mult:11.1f}x"
f" {'' if good else ' <-- did not reduce loss'}")
print(f"[{'PASS' if ok else 'FAIL'}] all optimizers reduce the loss")
return ok
def check_lion_sign_update():
"""Lion's step must be exactly +/- lr per non-zero coordinate."""
p = torch.nn.Parameter(torch.zeros(4))
o = Lion([p], lr=0.1, betas=(0.9, 0.99), weight_decay=0.0)
p.grad = torch.tensor([1.0, -1.0, 2.0, 0.0])
o.step()
expect = torch.tensor([-0.1, 0.1, -0.1, 0.0])
ok = torch.allclose(p.data, expect, atol=1e-6)
# The stored momentum must use beta2, not the beta1 used for the direction.
m_ok = torch.allclose(o.state[p]["m"],
0.01 * torch.tensor([1.0, -1.0, 2.0, 0.0]), atol=1e-6)
print(f"[{'PASS' if ok and m_ok else 'FAIL'}] Lion: step is +/-lr per "
f"coordinate and momentum uses the second beta")
return ok and m_ok
def check_mars_clip():
"""MARS clips the corrected gradient to unit norm."""
p = torch.nn.Parameter(torch.zeros(3))
o = MARS([p], lr=0.1, gamma=0.025, weight_decay=0.0, max_norm=1.0)
g = torch.tensor([1.0, 2.0, 3.0])
p.grad = g.clone()
o.step()
stored = torch.allclose(o.state[p]["g_prev"], g)
# After clipping, the first moment can be no larger than (1-b1) in norm.
b1 = o.param_groups[0]["betas"][0]
ok = o.state[p]["m"].norm().item() <= (1 - b1) + 1e-6
print(f"[{'PASS' if ok and stored else 'FAIL'}] MARS: corrected gradient "
f"clipped to unit norm, previous gradient cached")
return ok and stored
def check_newton_schulz():
"""Muon's iteration must compress the singular spectrum towards 1."""
from .optimizers.muon import newton_schulz
torch.manual_seed(0)
ok = True
print()
for shape in [(512, 512), (2048, 512), (512, 2048)]:
G = torch.randn(*shape)
O = newton_schulz(G, steps=5)
ci = torch.linalg.svdvals(G.float())
co = torch.linalg.svdvals(O.float())
before, after = (ci.max() / ci.min()).item(), (co.max() / co.min()).item()
good = after < before
ok &= good
print(f" {str(shape):>12s} condition number {before:8.2f} -> "
f"{after:6.2f}, mean singular value {co.mean():.3f}")
print(f"[{'PASS' if ok else 'FAIL'}] Muon: Newton-Schulz orthogonalises")
return ok
def check_sophia_scale():
"""Sophia's clip must not saturate once the batch scaling is applied.
With bs_scale left at 1 the denominator is orders of magnitude too small,
every coordinate saturates the clip, and Sophia silently becomes sign-SGD.
"""
def saturation(bs_scale):
torch.manual_seed(0)
m = _model()
o = SophiaG(m.parameters(), lr=1e-3, rho=0.05)
torch.manual_seed(1)
for i in range(60):
x, y = torch.randn(16, 16), torch.randint(0, 4, (16,))
logits = m(x)
o.zero_grad()
F.cross_entropy(logits, y).backward()
if (i + 1) % 10 == 0:
samp = torch.multinomial(logits.detach().softmax(-1), 1).squeeze(-1)
o.zero_grad()
F.cross_entropy(m(x), samp).backward()
o.update_hessian()
o.zero_grad()
F.cross_entropy(m(x), y).backward()
o.step(bs_scale=bs_scale)
tot = cl = 0
for p in m.parameters():
st = o.state[p]
r = st["m"] / (0.05 * bs_scale * st["h"]).clamp_min(1e-12)
tot += r.numel()
cl += (r.abs() >= 1).sum().item()
return cl / tot
lo, hi = saturation(1.0), saturation(4096.0)
ok = hi < lo
print(f"\n[{'PASS' if ok else 'FAIL'}] Sophia: clip saturation "
f"{lo:.3f} at bs_scale=1 (degenerates to sign-SGD) -> "
f"{hi:.3f} at bs_scale=4096")
return ok
def main():
results = [
check_adamw_matches_torch(),
check_lion_sign_update(),
check_mars_clip(),
check_newton_schulz(),
check_sophia_scale(),
check_all_reduce_loss(),
]
print(f"\n{sum(results)}/{len(results)} checks passed")
return 0 if all(results) else 1
if __name__ == "__main__":
raise SystemExit(main())
|