Step 6 · The training loop that converges
The model is the easy part. This module is the one that decides whether your run works or wastes a week.
In 60 seconds
Step 6 · The training loop that converges
The model is the easy part. This module is the one that decides whether your run works or wastes a week.
import math, torch
model = TinyGPT(Config()).to(device)
model = torch.compile(model) # big speedup, one line
# weight decay on matrices, none on norms and biases
decay = [p for n, p in model.named_parameters() if p.dim() >= 2]
nodecay = [p for n, p in model.named_parameters() if p.dim() < 2]
opt = torch.optim.AdamW(
[{"params": decay, "weight_decay": 0.1},
{"params": nodecay, "weight_decay": 0.0}],
lr=6e-4, betas=(0.9, 0.95), eps=1e-8, fused=True)
max_steps, warmup, min_lr = 20000, 500, 6e-5
def lr_at(step):
if step < warmup: # linear warmup
return 6e-4 * (step + 1) / warmup
r = (step - warmup) / (max_steps - warmup) # cosine decay
return min_lr + 0.5 * (6e-4 - min_lr) * (1 + math.cos(math.pi * r))
scaler_dtype = torch.bfloat16
accum = 8 # gradient accumulation
for step in range(max_steps):
for g in opt.param_groups:
g["lr"] = lr_at(step)
opt.zero_grad(set_to_none=True)
for micro in range(accum): # simulate a big batch
x, y = get_batch("train")
with torch.autocast(device_type="cuda", dtype=scaler_dtype):
_, loss = model(x, y)
(loss / accum).backward()
norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
if step % 100 == 0:
print(f"step {step} loss {loss.item():.3f} lr {lr_at(step):.2e} gnorm {norm:.2f}")Every line that matters, and why
| Choice | Typical value | What goes wrong without it |
|---|---|---|
| AdamW, not Adam | weight_decay 0.1 on matrices only | Decaying norms and biases hurts; this split is standard |
| beta2 = 0.95 | not the 0.999 default | 0.999 adapts too slowly and destabilises large models |
| Warmup | 200–2000 steps | Early huge updates on random weights blow the run up in the first minute |
| Cosine decay to ~10% | over the whole run | Constant LR plateaus well above the achievable loss |
| Gradient clipping | norm 1.0 | A single bad batch produces a spike that never recovers |
| Gradient accumulation | to reach 0.5–4M tokens per step | Small batches give noisy gradients and worse final loss |
| bfloat16 autocast | not fp16 | fp16 needs loss scaling and still overflows; bf16 has the range of fp32 |
Watch these four numbers, not just the loss
- 1
Training loss
Should fall fast then slowly. A flat line from step zero means a data bug. A sudden spike means a bad batch or too-high learning rate. - 2
Validation loss
Held-out data. When it stops falling while training loss keeps falling, you are memorising. For large pretraining runs on fresh data this rarely happens; for fine-tuning it happens in minutes. - 3
Gradient norm
Should be stable, roughly in the range 0.1–1.0. Growing steadily means instability building; frequent clipping means your learning rate is too high. - 4
Tokens per second
Your real currency. Everything in Module 46 is about this number, and it decides whether the run takes three days or three weeks.
Checkpoint like you expect to crash
def save(step):
torch.save({
"step": step,
"model": model.state_dict(),
"opt": opt.state_dict(), # optimiser state is essential
"config": cfg.__dict__,
"rng": torch.get_rng_state(),
}, f"ckpt_{step}.pt")Watch and read more
Lab
A training run you deliberately broke four ways.
The problem
print(f"step {step} loss {loss.item():.3f} gnorm {norm:.2f} lr {lr:.2e}")
# Learn to read these three numbers together — they diagnose almost everything.You are done when
Hard questions
Try to answer before you reveal. If you can answer these, you understood the lesson.
Q1Your loss spikes at step 8,000 and never recovers. You have checkpoints every 1,000 steps. Write the exact recovery procedure.Reveal
Questions people ask
How do I choose a learning rate?
Start from a known-good value for your scale — around 6e-4 for a small model, falling to roughly 1.5e-4 at 7B and lower still at frontier scale. Larger models need smaller learning rates. If you must search, run 200 steps at several values and pick the largest that stays stable.
What batch size?
In tokens, not sequences. Small models do well around 0.5M tokens per step; large pretraining runs use 4M to 16M. Reach it with gradient accumulation if your GPUs cannot hold it in one go.
My loss went to NaN. What now?
Lower the learning rate, confirm warmup is active, check gradient clipping is on, and confirm you are using bf16 rather than fp16. If it still happens, look for a corrupt data shard — one file of garbage bytes will do it.
Should I use torch.compile?
Yes. It is typically a 1.3–2x speedup for one line, and the compile time is paid once. Disable it while debugging, because the error messages get much harder to read.
Lesson test
5 questions. Get 3 right (60%) to pass and complete this lesson.
Sign in with your phone number to take the test and save your progress