rzen
234a37f8fd
- src/train.py: AdamW (decay on matrices only), cosine LR w/ warmup, grad clipping, periodic train/val loss, CSV logging, resumable checkpoints. --overfit mode runs the overfit-one-batch sanity check. MPS, fp32. - configs/v1.py: TrainConfig (batch 64, peak LR 6e-4, 20k iters, etc.) - tests/test_train.py: LR schedule + optimizer grouping (3 tests, no run)
73 lines
2.1 KiB
Python
73 lines
2.1 KiB
Python
"""Tests for training helpers that don't require an optimization run.
|
|
|
|
Covers the LR schedule (a pure function) and the optimizer param grouping.
|
|
Run: `.venv/bin/python tests/test_train.py`
|
|
"""
|
|
|
|
import sys
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent
|
|
sys.path.insert(0, str(ROOT / "src"))
|
|
sys.path.insert(0, str(ROOT / "configs"))
|
|
|
|
from model import GPT # noqa: E402
|
|
from train import configure_optimizer, cosine_lr # noqa: E402
|
|
from v1 import ModelConfig # noqa: E402
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TC:
|
|
learning_rate: float = 1e-3
|
|
min_lr: float = 1e-4
|
|
warmup_iters: int = 100
|
|
max_iters: int = 1000
|
|
weight_decay: float = 0.1
|
|
beta1: float = 0.9
|
|
beta2: float = 0.95
|
|
|
|
|
|
def test_lr_warmup_is_linear():
|
|
tc = TC()
|
|
# Ramps from ~0 up to the peak across warmup_iters.
|
|
assert cosine_lr(0, tc) < cosine_lr(50, tc) < cosine_lr(99, tc)
|
|
assert abs(cosine_lr(tc.warmup_iters - 1, tc) - tc.learning_rate) < 1e-9
|
|
|
|
|
|
def test_lr_peaks_then_decays_to_floor():
|
|
tc = TC()
|
|
peak = cosine_lr(tc.warmup_iters, tc)
|
|
assert abs(peak - tc.learning_rate) < 1e-6
|
|
# Monotonically decreasing through the cosine phase.
|
|
mid = cosine_lr(tc.warmup_iters + (tc.max_iters - tc.warmup_iters) // 2, tc)
|
|
assert tc.min_lr < mid < peak
|
|
# Bottoms out at the floor.
|
|
assert abs(cosine_lr(tc.max_iters, tc) - tc.min_lr) < 1e-9
|
|
assert abs(cosine_lr(tc.max_iters + 500, tc) - tc.min_lr) < 1e-9
|
|
|
|
|
|
def test_optimizer_grouping():
|
|
model = GPT(ModelConfig())
|
|
opt = configure_optimizer(model, TC())
|
|
decay_group, no_decay_group = opt.param_groups
|
|
assert decay_group["weight_decay"] == 0.1
|
|
assert no_decay_group["weight_decay"] == 0.0
|
|
# Norm/bias params (1-D) must be in the no-decay group only.
|
|
assert all(p.dim() >= 2 for p in decay_group["params"])
|
|
assert all(p.dim() < 2 for p in no_decay_group["params"])
|
|
|
|
|
|
def main() -> None:
|
|
tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
|
|
for t in tests:
|
|
t()
|
|
print(f"ok {t.__name__}")
|
|
print(f"\n{len(tests)} passed")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|