Phase 4: training loop (code only — for user to run)

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

tests/test_train.py

@@ -0,0 +1,72 @@
+"""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()