Upload folder using huggingface_hub

5407bdc verified about 1 month ago

1.96 kB

"""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"]

	"""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"]