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)
This commit is contained in:
+26
-1
@@ -21,5 +21,30 @@ class ModelConfig:
|
||||
d_ff: int = 1024 # MLP hidden width (4 x d_model)
|
||||
|
||||
|
||||
# The single source of truth other modules import.
|
||||
@dataclass(frozen=True)
|
||||
class TrainConfig:
|
||||
# Data / batching. tokens per step = batch_size * context_length = 16,384.
|
||||
batch_size: int = 64
|
||||
|
||||
# AdamW optimiser.
|
||||
learning_rate: float = 6e-4 # peak LR (reached after warmup)
|
||||
min_lr: float = 6e-5 # cosine-decay floor (~10% of peak)
|
||||
warmup_iters: int = 200
|
||||
max_iters: int = 20_000 # ~327M tokens (~0.6 epoch of the 555M corpus)
|
||||
weight_decay: float = 0.1 # applied to matrices only, not norms/biases
|
||||
beta1: float = 0.9
|
||||
beta2: float = 0.95
|
||||
grad_clip: float = 1.0 # clip global grad norm to this
|
||||
|
||||
# Evaluation, logging, checkpointing (all in optimiser steps).
|
||||
eval_interval: int = 500 # how often to measure train/val loss
|
||||
eval_iters: int = 100 # batches averaged per loss estimate
|
||||
log_interval: int = 20 # how often to print step/loss/lr/speed
|
||||
checkpoint_interval: int = 1000
|
||||
|
||||
seed: int = 1337
|
||||
|
||||
|
||||
# Single source of truth other modules import.
|
||||
model = ModelConfig()
|
||||
train = TrainConfig()
|
||||
|
||||
+201
@@ -0,0 +1,201 @@
|
||||
"""Training loop: AdamW, cosine LR schedule with warmup, gradient clipping,
|
||||
periodic validation, CSV loss logging, and resumable checkpoints.
|
||||
|
||||
Run on MPS (Apple GPU). Everything is fp32 for clarity — mixed precision is an
|
||||
iteration-2 speedup.
|
||||
|
||||
Usage:
|
||||
python src/train.py --config v1 # train from scratch
|
||||
python src/train.py --config v1 --resume # continue from last checkpoint
|
||||
python src/train.py --config v1 --overfit # sanity check: overfit one batch
|
||||
|
||||
The overfit run is the recommended first step: it trains on a single fixed
|
||||
batch and the loss should collapse toward ~0 within a few hundred steps. If it
|
||||
does, the model / loss / optimizer are wired correctly and a real run is worth
|
||||
starting.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import importlib
|
||||
import math
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import asdict
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
ROOT = Path(__file__).resolve().parent.parent
|
||||
sys.path.insert(0, str(ROOT / "src"))
|
||||
sys.path.insert(0, str(ROOT / "configs"))
|
||||
|
||||
from dataset import get_batch, load_tokens # noqa: E402
|
||||
from model import GPT # noqa: E402
|
||||
|
||||
|
||||
def pick_device() -> str:
|
||||
if torch.backends.mps.is_available():
|
||||
return "mps"
|
||||
return "cpu"
|
||||
|
||||
|
||||
def cosine_lr(it: int, tc) -> float:
|
||||
"""Linear warmup, then cosine decay from peak LR down to the floor."""
|
||||
if it < tc.warmup_iters:
|
||||
return tc.learning_rate * (it + 1) / tc.warmup_iters
|
||||
if it >= tc.max_iters:
|
||||
return tc.min_lr
|
||||
progress = (it - tc.warmup_iters) / (tc.max_iters - tc.warmup_iters)
|
||||
coeff = 0.5 * (1.0 + math.cos(math.pi * progress))
|
||||
return tc.min_lr + coeff * (tc.learning_rate - tc.min_lr)
|
||||
|
||||
|
||||
def configure_optimizer(model: GPT, tc) -> torch.optim.Optimizer:
|
||||
"""AdamW with weight decay on matrices (dim >= 2) but not on norms/biases."""
|
||||
decay = [p for p in model.parameters() if p.dim() >= 2]
|
||||
no_decay = [p for p in model.parameters() if p.dim() < 2]
|
||||
groups = [
|
||||
{"params": decay, "weight_decay": tc.weight_decay},
|
||||
{"params": no_decay, "weight_decay": 0.0},
|
||||
]
|
||||
return torch.optim.AdamW(groups, lr=tc.learning_rate, betas=(tc.beta1, tc.beta2))
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def estimate_loss(model, splits, tc, ctx_len, device) -> dict[str, float]:
|
||||
"""Average loss over eval_iters random batches for each split."""
|
||||
model.eval()
|
||||
out = {}
|
||||
for name, data in splits.items():
|
||||
losses = torch.zeros(tc.eval_iters)
|
||||
for i in range(tc.eval_iters):
|
||||
x, y = get_batch(data, tc.batch_size, ctx_len, device)
|
||||
_, loss = model(x, y)
|
||||
losses[i] = loss.item()
|
||||
out[name] = losses.mean().item()
|
||||
model.train()
|
||||
return out
|
||||
|
||||
|
||||
def overfit_one_batch(model, opt, train_data, tc, ctx_len, device, steps=400):
|
||||
"""Sanity check: repeatedly fit a single batch; loss should approach 0."""
|
||||
gen = torch.Generator().manual_seed(0)
|
||||
x, y = get_batch(train_data, tc.batch_size, ctx_len, device, generator=gen)
|
||||
print(f"overfitting one batch ({tc.batch_size}x{ctx_len}) for {steps} steps")
|
||||
for it in range(steps):
|
||||
_, loss = model(x, y)
|
||||
opt.zero_grad(set_to_none=True)
|
||||
loss.backward()
|
||||
opt.step()
|
||||
if it % 20 == 0 or it == steps - 1:
|
||||
print(f" step {it:4d} loss {loss.item():.4f}")
|
||||
final = loss.item()
|
||||
print(f"final loss {final:.4f} — "
|
||||
+ ("OK, wiring looks correct" if final < 0.5 else "TOO HIGH, investigate"))
|
||||
|
||||
|
||||
def main() -> None:
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--config", default="v1", help="config module in configs/")
|
||||
ap.add_argument("--resume", action="store_true")
|
||||
ap.add_argument("--overfit", action="store_true")
|
||||
args = ap.parse_args()
|
||||
|
||||
cfg = importlib.import_module(args.config)
|
||||
mc, tc = cfg.model, cfg.train
|
||||
ctx_len = mc.context_length
|
||||
device = pick_device()
|
||||
torch.manual_seed(tc.seed)
|
||||
print(f"device: {device}")
|
||||
|
||||
# Data (memory-mapped; never fully loaded).
|
||||
tokens_dir = ROOT / "data" / "tokens"
|
||||
train_data = load_tokens(tokens_dir / "train.bin")
|
||||
val_data = load_tokens(tokens_dir / "val.bin")
|
||||
print(f"train tokens: {len(train_data):,} | val tokens: {len(val_data):,}")
|
||||
|
||||
model = GPT(mc).to(device)
|
||||
print(f"model params: {model.num_params():,}")
|
||||
opt = configure_optimizer(model, tc)
|
||||
|
||||
if args.overfit:
|
||||
overfit_one_batch(model, opt, train_data, tc, ctx_len, device)
|
||||
return
|
||||
|
||||
ckpt_dir = ROOT / "checkpoints" / args.config
|
||||
ckpt_dir.mkdir(parents=True, exist_ok=True)
|
||||
ckpt_path = ckpt_dir / "ckpt.pt"
|
||||
reports_dir = ROOT / "reports"
|
||||
reports_dir.mkdir(exist_ok=True)
|
||||
csv_path = reports_dir / f"loss-{args.config}.csv"
|
||||
|
||||
start_iter = 0
|
||||
best_val = float("inf")
|
||||
if args.resume and ckpt_path.exists():
|
||||
ckpt = torch.load(ckpt_path, map_location=device)
|
||||
model.load_state_dict(ckpt["model"])
|
||||
opt.load_state_dict(ckpt["optimizer"])
|
||||
start_iter = ckpt["iter"] + 1
|
||||
best_val = ckpt.get("best_val", best_val)
|
||||
print(f"resumed from iter {start_iter}")
|
||||
else:
|
||||
# Fresh run: start the CSV with a header.
|
||||
with open(csv_path, "w", newline="") as f:
|
||||
csv.writer(f).writerow(["iter", "train_loss", "val_loss", "lr"])
|
||||
|
||||
model.train()
|
||||
t0 = time.perf_counter()
|
||||
for it in range(start_iter, tc.max_iters):
|
||||
lr = cosine_lr(it, tc)
|
||||
for group in opt.param_groups:
|
||||
group["lr"] = lr
|
||||
|
||||
x, y = get_batch(train_data, tc.batch_size, ctx_len, device)
|
||||
_, loss = model(x, y)
|
||||
opt.zero_grad(set_to_none=True)
|
||||
loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), tc.grad_clip)
|
||||
opt.step()
|
||||
|
||||
if it % tc.log_interval == 0:
|
||||
dt = time.perf_counter() - t0
|
||||
ips = (it - start_iter + 1) / dt
|
||||
eta_min = (tc.max_iters - it) / ips / 60 if ips > 0 else 0
|
||||
print(f"iter {it:6d}/{tc.max_iters} loss {loss.item():.4f} "
|
||||
f"lr {lr:.2e} {ips:5.1f} it/s eta {eta_min:5.1f}m")
|
||||
|
||||
if it > 0 and it % tc.eval_interval == 0:
|
||||
stats = estimate_loss(model, {"train": train_data, "val": val_data},
|
||||
tc, ctx_len, device)
|
||||
print(f" eval train {stats['train']:.4f} val {stats['val']:.4f}")
|
||||
with open(csv_path, "a", newline="") as f:
|
||||
csv.writer(f).writerow([it, f"{stats['train']:.4f}",
|
||||
f"{stats['val']:.4f}", f"{lr:.6f}"])
|
||||
best_val = min(best_val, stats["val"])
|
||||
|
||||
if it > 0 and it % tc.checkpoint_interval == 0:
|
||||
torch.save({
|
||||
"model": model.state_dict(),
|
||||
"optimizer": opt.state_dict(),
|
||||
"iter": it,
|
||||
"best_val": best_val,
|
||||
"model_cfg": asdict(mc),
|
||||
}, ckpt_path)
|
||||
|
||||
# Final checkpoint.
|
||||
torch.save({
|
||||
"model": model.state_dict(),
|
||||
"optimizer": opt.state_dict(),
|
||||
"iter": tc.max_iters - 1,
|
||||
"best_val": best_val,
|
||||
"model_cfg": asdict(mc),
|
||||
}, ckpt_path)
|
||||
print(f"done. checkpoint: {ckpt_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user