"""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())