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