Part IV — Pretraining

The training loop — AdamW, batches, curves

Lesson 09 turned the model's error into one number — cross-entropy, the mean negative log-likelihood of the true next token over a batch. That number is a snapshot: it says how wrong we are right now, on a handful of examples. This lesson turns the snapshot into a movie. We wrap the forward pass, the loss, backpropagation, an optimizer step, and a gradient reset into a loop that runs over every batch in the corpus, epoch after epoch — and we learn to read the two curves it produces, because one of them will quietly betray us on a small dataset.

What the last step left broken

A single cross-entropy value on a single batch is a thermometer reading, not a treatment. Nothing in lesson 09 actually changed the weights — we measured the error and stopped. To learn, we have to (a) compute the gradient of that loss with respect to every parameter, (b) nudge the parameters downhill, and (c) do it again on the next batch, and the next, until the whole corpus has taught the model something. And we need a way to tell "the model is genuinely learning" apart from "the model is memorizing these exact sentences" — a distinction one loss number cannot make.

Linear position

Forced by: lesson 09 gave one loss on one batch, but a lone number changes nothing — training means lowering that number over an entire corpus, across many batches and many passes, without the optimization blowing up.
New idea: the training loop — for each batch: forward → loss → backprop → optimizer step → zero the gradients; repeat over epochs. The optimizer is AdamW (a per-parameter adaptive step with decoupled weight decay), and its learning rate is the master knob — too high and the loss diverges, too low and it crawls. Split the data into train and validation and watch both curves: when the validation curve flattens or turns up while the training curve keeps falling, you are overfitting.
Forces next: training a real model from scratch is expensive, and the tiny toy model we can afford is dull — greedy decoding makes it repeat itself. Lesson 11 adds sampling control (temperature, top-k) and loads OpenAI's published GPT-2 weights so we skip the bill.

The plan

Four moves. (1) Write the loop: the five steps that repeat for every batch, and what an epoch is. (2) Pick the optimizer: why plain gradient descent gives way to AdamW, why the learning rate dominates everything, and where warmup and cosine schedules fit. (3) Split train from validation and read the loss curve — the shape that means learning, and the fork that means overfitting. (4) Decide what to watch and when to stop. Then drive a deterministic training-loop simulator and make the loss diverge, overfit, and converge on demand.

1 · The loop: five steps, repeated

All of training is a short loop wrapped in a longer loop. The inner loop walks through the corpus one batch at a time; the outer loop repeats that walk once per epoch. One epoch is one complete pass over the training set, and the number of batches in an epoch is just the training-set size divided by the batch size. Inside the inner loop, exactly five things happen, in this order:

step 1

Forward

Run the batch of token windows through the GPT to get logits, shape (B, T, V).

step 2

Loss

Cross-entropy of the logits against the shifted targets — the one number from lesson 09.

step 3

Backward

Backprop: compute ∂loss/∂θ for every parameter θ at once.

step 4

Step

The optimizer nudges each parameter a little way downhill using its gradient.

step 5

Zero grads → repeat

Reset the accumulated gradients to zero, then grab the next batch and go again.

In PyTorch-flavoured pseudo-code — paraphrased, not the book's listing — the whole thing is startlingly small:

for epoch in range(num_epochs):
    for input_batch, target_batch in train_loader:   # inner loop over batches
        optimizer.zero_grad()                        # step 5, done first so grads don't accumulate
        logits = model(input_batch)                  # step 1  forward  -> (B, T, V)
        loss   = cross_entropy(logits, target_batch) # step 2  the scalar from lesson 09
        loss.backward()                              # step 3  fill every .grad
        optimizer.step()                             # step 4  apply the update
    # end of epoch: measure val loss, print a sample generation

Two subtleties earn their keep. First, why zero the gradients at all? Autograd frameworks accumulate gradients by default — each backward() adds into the existing .grad buffers rather than replacing them. If you forget zero_grad(), batch 2's update is contaminated by batch 1's stale gradient, batch 3 by both, and the effective step size balloons until training destabilizes. It is the single most common training bug, and it produces a loss curve that looks eerily like "the learning rate is too high." Second, notice that the loop only touches the training set. To know whether the model is learning language rather than memorizing sentences, we periodically run a second, gradient-free pass over a held-out validation set — the source of the second curve we study in section 3.

2 · The optimizer: from gradient descent to AdamW

Plain gradient descent, and why it isn't enough

Backprop hands us a gradient — the direction in parameter space that increases the loss fastest. To decrease the loss we step the opposite way. The simplest rule is vanilla gradient descent: subtract a small multiple of the gradient from each parameter.

θ ← θ − η · ∂loss/∂θ

Here η (eta) is the learning rate: the size of the step. This one line trains, but slowly and fragilely, for two reasons. Different parameters want different step sizes — a weight whose gradient is consistently tiny should move faster, one whose gradient thrashes should move more cautiously — yet a single global η treats them identically. And a raw gradient is noisy from batch to batch, so the path down the loss surface zig-zags.

AdamW: an adaptive, smoothed step per parameter

Adam fixes both problems by keeping two running averages per parameter: a smoothed gradient (first moment, the "momentum" — which direction have I been heading?) and a smoothed squared gradient (second moment — how large and how volatile have my gradients been?). Each parameter's step is the momentum divided by the square root of that volatility. The effect is that every parameter gets its own effective learning rate, automatically: chronically small-gradient weights speed up, thrashing weights are damped. This is why Adam-family optimizers train deep networks far more reliably than plain descent.

The W in AdamW is decoupled weight decay. Weight decay gently pulls every parameter toward zero each step, a regularizer that discourages the model from leaning on any one large weight — which curbs overfitting. Ordinary Adam folds that pull into the gradient, where the adaptive scaling distorts it; AdamW applies the decay as a separate, clean shrink after the gradient step, so the regularization behaves as intended. For that reason AdamW is the default optimizer for training LLMs, and it is what we use here. A typical setup mirrors the book's: a learning rate around 4×10⁻⁴ and a weight decay around 0.1.

The learning rate is the master knob

Of all the dials on the optimizer, the learning rate matters most, and it has a narrow good range flanked by two failure modes.

learning rate	what happens	curve signature
too high	each step overshoots the valley and lands higher up the far wall; errors compound.	loss shoots up, often to NaN — divergence.
good	steps are large enough to make progress, small enough to stay in the valley.	loss falls smoothly, then eases toward a floor.
too low	every step is a shuffle; the model does learn, but glacially.	loss inches down; you run out of compute before you finish.

Because both edges are painful, practitioners rarely hold the learning rate fixed. Two schedules are near-universal and the widget's fixed rate is a stand-in for their average. Warmup ramps η up from near-zero over the first few hundred steps: early gradients are wild, and a full-size step then can blow the run apart before it starts. Cosine decay then eases η back down over the rest of training, taking big strides early and fine, careful ones near the end. We keep the schedule story brief here on purpose — the derivations, gradient clipping, and the stability tricks that make very large runs survive live in the systems track.

Where the schedule details live

Warmup, cosine annealing, gradient clipping, and mixed-precision stability are the meat of a scale-focused course. See cs336 · 04 Training for the AdamW update written out term by term, the learning-rate schedule, and the tricks that keep a billion-parameter run from diverging.

3 · Train vs validation, and the shape of the curve

Why hold data back

The training loss only tells you how well the model fits the sentences it is actively being corrected on. A model with enough capacity can drive that number to nearly zero by memorizing — storing the training text verbatim rather than learning the patterns of language. To catch this, we carve the corpus into a training split (the model learns from it) and a validation split (the model never trains on it, we only measure loss on it). The gap between the two is the whole story:

Both fall together → the model is learning generalizable structure that transfers to unseen text. This is what you want.
Training keeps falling, validation flattens or turns up → the model is now fitting quirks of the training set that do not transfer. This is overfitting.

Overfitting is guaranteed on a tiny corpus

Follow the book's own run. It pretrains the GPT for ten epochs on a single short story — Edith Wharton's "The Verdict," about twenty thousand characters — and plots both losses. Both start near 9.9 and drop steeply for the first epoch or two. Then they split: the training loss slides on down toward a fraction of one, while the validation loss stalls around 6.4 and refuses to follow. Past roughly epoch 2 the two curves diverge for good.

The diagnosis, made concrete

The book confirms the memorization directly: search the model's generated text and you find whole phrases lifted verbatim from the story, like "quite insensible to the irony." With so few tokens and ten passes over them, the model has enough capacity to store the corpus outright — so the training loss collapses while the validation loss, measured on text it has never been corrected on, plateaus. This is not a bug in the loop; it is the expected outcome of training a large model many times over a tiny dataset. The realistic fix is more data (the book points to pretraining on tens of thousands of public-domain books, where the split does not appear), or fewer passes.

The lesson generalizes past this one story. The training curve measures fit; the validation curve measures generalization; and the moment they part is the moment additional training stops helping and starts hurting. On a large, diverse corpus the two stay close for a long time, which is exactly why "large" in the data sense is as important as "large" in the parameter sense (lesson 00).

4 · What you watch, and when to stop

During a run you keep three things in view, and together they tell you when to pull the plug.

The curves

Plot train and validation loss as training proceeds. Smooth descent = healthy. A sudden spike = the learning rate is too high (or you forgot zero_grad). The two curves splitting = overfitting.

Sample generations

Every epoch, generate a few tokens from a fixed prompt and read them. Early on it is comma-spam and repeated words; later it forms grammatical clauses. Qualitative, but it makes "loss went down" tangible.

The stopping point

Stop when the validation loss stops improving — training past that point only sharpens memorization. This is early stopping, and the validation curve is what triggers it.

There is one more thing to watch that this whole course keeps circling back to: cost. The book's ten-epoch toy run finishes in about five minutes on a laptop — because the model is small and the corpus is one story. A real pretraining run is many orders of magnitude larger in both, and the bill scales with the product of parameters and tokens processed. That expense is precisely why lesson 11 will load OpenAI's already-trained GPT-2 weights instead of paying to pretrain from scratch, and why the systems track builds its entire spine around a compute budget.

Where the cost accounting lives

The rule of thumb is that training compute is roughly C ≈ 6ND FLOPs for N parameters and D tokens, which is why "just train longer on more data" is a budget question, not a code question. See cs336 · 05 Accounting for the FLOP and memory arithmetic behind a training run.

5 · Drive the loop yourself

The widget below is a deterministic model of a training run — no randomness anywhere, just fixed formulas of the four sliders, so the same settings always draw the same two curves. It plots training loss (blue) and validation loss (orange) against optimizer steps, and reports a status of converged, diverged, or overfit. Three things are wired to fail exactly the way real training does. Push the learning rate past the stability edge and the curve blows up (divergence) — and note that a larger batch size gives smoother gradients, so it tolerates a higher rate before it breaks. Shrink the corpus size toward "one short story" and the validation curve peels off and turns upward while training marches to zero (overfitting), no matter how gently you tune the rate. Give it a sane rate and enough data and both curves settle onto a shared floor (convergence).

Training-loop simulator

Proves the two headline claims: the learning rate is the master knob (too high → the loss diverges), and a tiny corpus overfits no matter what (validation turns up while training → 0). Everything is a fixed deterministic function of the sliders — no random-number generator, no clock.

learning rate: 4.0e-4 batch size: 16 epochs: 10 corpus size: tiny · ~1 story

final train loss

—

final val loss

—

train–val gap

—

status

—

Failure modes & checklist

Failure modes

Forgetting zero_grad(). Gradients accumulate across batches; the effective step grows until the loss explodes. Signal: looks exactly like too-high learning rate — a loss that spikes upward.
Learning rate too high. Steps overshoot the valley; loss climbs, often to NaN. Signal: the curve turns up early and never recovers.
Learning rate too low. The model learns, but so slowly you run out of budget. Signal: loss barely moves over many steps.
Ignoring the validation curve. Chasing training loss to zero on a small corpus just memorizes it. Signal: train and val split; val flattens or rises.
No warmup on a large run. A full-size first step on wild initial gradients blows the run apart. Signal: immediate divergence in the first dozen steps.

Checklist

The five steps in order: forward → loss → backward → step → zero_grad.
AdamW as the optimizer, with a small weight decay (≈ 0.1).
A learning rate in the good range, ideally with warmup + cosine decay.
A held-out validation split you never train on.
Both curves logged; sample generations printed each epoch.
Stop when validation loss stops improving (early stopping).

Checkpoint

Try it

Start with the defaults and confirm the status reads overfit — the corpus slider is near "tiny," so the orange validation curve peels away and turns up while blue marches down, and the train–val gap is large. Now drag corpus size all the way to the right and watch the two curves collapse back together onto a shared floor: status flips to converged. Return corpus to the middle, then drag learning rate up past roughly 1e-3: the curve turns upward and the status flips to diverged. Finally, from that diverged state, push batch size to the right — the smoother gradient buys back some stability and the run can recover. Which of the three properties from lesson 00 — accurate, trainable, steerable — is this entire loop in service of?

Where this points next

You can now train the model: the loop lowers cross-entropy over the corpus, AdamW makes each step adaptive and stable, and the two curves tell you whether you are learning or memorizing. But two problems remain, and together they point straight at the next lesson. First, the toy model you can actually afford to train from scratch is weak, and when you generate from it with the greedy "always take the top token" rule, it falls into short repeating loops — dull and unusable. Second, training a model strong enough to be interesting costs far more compute than the five-minute laptop run. Lesson 11, Decoding & loading pretrained weights, solves both: temperature and top-k sampling to control the randomness of generation, and a way to load OpenAI's published GPT-2 weights directly into the architecture we built — a strong base model, for free.

Takeaway

Training is a loop: for every batch, run the forward pass, compute the cross-entropy loss (09), backpropagate to get every gradient, take an optimizer step, and zero the gradients — repeated over batches and epochs. The optimizer is AdamW: per-parameter adaptive steps (from smoothed first and second moments) plus decoupled weight decay, and its learning rate is the master knob — too high diverges, too low crawls, with warmup + cosine bridging the two. Split the data into train and validation and read both curves: when validation flattens or turns up while training keeps falling, you are overfitting — guaranteed on a tiny corpus, and the cue to stop or get more data.

Interview prompts

List the steps of the inner training loop in order. (§1 — forward → cross-entropy loss → backward (backprop) → optimizer step → zero gradients; repeat over batches, then over epochs.)
Why must you call zero_grad() each iteration, and what happens if you don't? (§1 — autograd accumulates gradients into .grad by default; skipping the reset lets stale gradients pile up so the effective step grows and the loss destabilizes — it mimics too-high learning rate.)
What does AdamW add over plain gradient descent, and what is the "W"? (§2 — per-parameter adaptive step sizes from smoothed first and second moments (momentum + variance), plus decoupled weight decay ("W") applied separately from the gradient for cleaner regularization.)
Why is the learning rate called the master knob? (§2 — too high overshoots the valley and diverges (often to NaN); too low crawls and wastes compute; only a narrow middle range trains well, which is why schedules like warmup + cosine are used.)
What is the difference between the training and validation loss, and what does a growing gap mean? (§3 — training loss measures fit on data the model is corrected on; validation loss measures generalization on held-out data; a widening gap (val flat/rising, train falling) means overfitting.)
Why does a tiny corpus overfit no matter how you tune the optimizer? (§3 — a high-capacity model trained for several passes over few tokens can memorize them verbatim, collapsing train loss while val loss plateaus; the fix is more data or fewer epochs, not a different learning rate.)
What is early stopping and which signal triggers it? (§4 — halting when the validation loss stops improving, because further training only deepens memorization; the validation curve is the trigger.)

Companion reads: cs336 · 04 Training writes out the AdamW update, the learning-rate schedule, and large-run stability; cs336 · 05 Accounting works the compute and memory cost of a run; gpt_mini · 02 Pretrain runs this loop on an assumed architecture.