File size: 1,957 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 | """From-scratch optimizers, one per category of Table 1 in
*Fantastic Pretraining Optimizers and Where to Find Them* (Wen et al., 2025).
category implementation state / param
-------------------------------- ------------------ -------------
1. AdamW AdamW 2x
2. Variance-reduced AdamW variant MARS 3x
3. Memory-efficient Lion 1x
4. Matrix-based Muon 1x (2D) / 2x
5. Hessian-based (bonus) Sophia-G 2x
``torch.optim.Optimizer`` is used only as a base class; no ``torch.optim``
update rule is called anywhere.
"""
from .adamw import AdamW
from .lion import Lion
from .mars import MARS
from .muon import Muon
from .sophia import SophiaG
OPTIMIZERS = {
"adamw": AdamW,
"mars": MARS,
"lion": Lion,
"muon": Muon,
"sophia": SophiaG,
}
CATEGORY = {
"adamw": "1. AdamW",
"mars": "2. Variance-reduced AdamW variant",
"lion": "3. Memory-efficient",
"muon": "4. Matrix-based",
"sophia": "5. Hessian-based (bonus)",
}
# Peak learning rates are the values selected by src/part2/sweep.py; the
# remaining hyperparameters are each method's published defaults and were not
# swept. Note how far apart the selected rates are -- Muon tolerates 500x
# Sophia's -- which is exactly why every optimizer needed its own grid.
DEFAULT_HPARAMS = {
"adamw": dict(lr=1e-3, betas=(0.9, 0.95), weight_decay=0.1),
"mars": dict(lr=3e-3, betas=(0.95, 0.99), weight_decay=0.1, gamma=0.025),
"lion": dict(lr=3e-4, betas=(0.9, 0.99), weight_decay=0.6),
"muon": dict(lr=5e-2, momentum=0.95, weight_decay=0.1, adamw_lr=1e-3),
"sophia": dict(lr=1e-4, betas=(0.965, 0.99), rho=0.05, weight_decay=0.1),
}
__all__ = ["AdamW", "MARS", "Lion", "Muon", "SophiaG",
"OPTIMIZERS", "CATEGORY", "DEFAULT_HPARAMS"]
|