Chapter 6 · Section 6.3 · 45 min

Training the model

The model from §6.2 knows nothing. Its loss is $\ln V$ , the cost of a uniform guess. Training it looks like one line, loss.backward(); opt.step(), and for a toy model that line mostly works. Yet nearly every number in the configuration of a modern GPT-2 reproduction (most of them borrowed from GPT-3) exists because the one-line version fails in some specific way at scale. There's the $\beta_2$ of 0.95 instead of Adam's default 0.999, the 700 steps of warmup, the clipping threshold of 1.0, and the rule that only matrices get weight decay. You built each of those pieces in Chapter 0, AdamW, the schedule, and clipping in §0.6, and accumulation and mixed precision in §0.9. This section does the part that is new for a GPT. It writes the loss over a $(B, T, V)$ block of logits, assembles the pieces into one correct optimizer step, trains our GPT on a CPU in about two seconds, and reads what the loss curves say.

Keep one picture in mind. A training run is a single curve, loss against steps, and that curve is the main instrument you'll have when a run costs real money. A doctor reads a chart. By the end of this section you'll read a loss curve and say "learning rate too high", "overfitting", or "that's a data bug" from its shape alone. Here's one to start with. The loss falls quickly, flattens at 2.3, then suddenly jumps to 4 and stays there. What happened? Keep a guess in mind until the section on loss curves.

The data, again

Each section of this chapter runs on its own, so the first cell rebuilds the pipeline of §6.1 with the same eight sonnets.

Python · cell 1

The loss, cross-entropy over the vocabulary

The model outputs logits $\mathbf{Z} \in \mathbb{R}^{B \times T \times V}$ . Flatten the batch and time axes into $N = BT$ rows, so row $n$ has logits $\mathbf{z}_n \in \mathbb{R}^V$ and target $y_n$ . Some targets should not count, such as padding, or in fine-tuning the prompt tokens (§14.3). We mark those with a special value, conventionally $-100$ , the ignore index. With $\mathcal{C}$ the set of rows whose target is real (the counted rows), the loss is

\mathcal{L} = \frac{1}{|\mathcal{C}|}\sum_{n \in \mathcal{C}} \Big(\log \sum_{i=1}^{V} e^{z_{n,i}} - z_{n, y_n}\Big),

the average negative log-probability of the correct next token (§0.5 explains why cross-entropy is the right loss for a distribution). Its gradient is the familiar "predicted minus actual" of §0.2, divided by the number of counted rows,

\frac{\partial \mathcal{L}}{\partial z_{n,i}} = \frac{\operatorname{softmax}(\mathbf{z}_n)_i - \mathbb{1}[i = y_n]}{|\mathcal{C}|} \quad \text{for } n \in \mathcal{C}, \qquad \text{and } 0 \text{ for ignored rows}.

In plain words, every counted position pushes up its target's logit and pushes down all the others in proportion to their probability. Ignored positions contribute nothing, not even to the denominator.

Two details decide whether this is right in practice. The log-sum-exp must subtract the row maximum before exponentiating (§0.4), because a logit of 100 overflows exp in float16 and float32 alike. And the mean must divide by the number of counted targets, not by $N$ . Divide by $N$ and a batch that is half padding gets half the gradient, so the effective learning rate depends on how much padding the batch happened to contain.

Python · cell 2

The second row was ignored, so the loss is just the first row's $-\log p$ , and the second row's gradient is all zeros. At GPT-2 scale the logits are also the biggest single tensor in the step. For $B = 16$ , $T = 1024$ , and $V = 50{,}257$ they hold 823 million numbers, 3.3 GB in float32, before the gradient of the same size. That's why production code computes the loss in chunks or with a fused kernel that never materializes all of them at once.

Interview

Cross-entropy over logits, with an ignore index

Implement ce_loss_and_grad(logits, targets, ignore_index=-100) for logits of shape (B, T, V) and integer targets of shape (B, T). Return the mean negative log-likelihood over non-ignored targets and its gradient with respect to the logits (same shape as the logits). It must be numerically stable for logits as large as 1000, ignored positions must get zero gradient, and the mean must be over the non-ignored count. You may assume at least one target is not ignored. Ignored targets may hold any value equal to ignore_index, which can be negative, so don't use them as indices.

import numpy as np

def ce_loss_and_grad(logits, targets, ignore_index=-100):
    """returns (loss: float, dlogits: same shape as logits)"""
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

The pieces from Chapter 0, in one step

A GPT training step uses five pieces you have already built. Each gets one line of recap here, with the setting a GPT-2 run uses, because the new work in this section is putting them together in the right order.

AdamW (§0.6) keeps running averages $\mathbf{m}_s$ and $\mathbf{v}_s$ of the gradient and its square, corrects them for starting at zero, steps along $\hat{\mathbf{m}}_s / (\sqrt{\hat{\mathbf{v}}_s} + \epsilon)$ , and shrinks each decayed weight by the factor $1 - \eta_s\lambda$ as a separate step. GPT-3's settings, which nanoGPT and llm.c reuse, are $\beta_1 = 0.9$ , $\beta_2 = 0.95$ , $\epsilon = 10^{-8}$ , and weight decay $\lambda = 0.1$ (Brown et al., 2020). §0.6 explains the short-memory $\beta_2$ .
Decay groups. Decay every tensor with two or more dimensions (the weight matrices and embeddings) and nothing else. A LayerNorm gain's natural value is 1 and a bias only shifts activations, so pulling either toward zero fights the model and buys no regularization.
Warmup, then cosine decay (§0.6) ramps the learning rate linearly to $\eta_{\max}$ over the first $S_w$ steps, then lowers it along half a cosine to $\eta_{\min}$ at the last step $S$ . The GPT-2 reproduction warms up for 700 steps to a peak of $6 \times 10^{-4}$ .
Clipping by global norm (§0.6) rescales the whole gradient to length $c = 1.0$ whenever it is longer, keeping its direction. Log the norm before clipping. It is the best early warning of a loss spike you have.
Gradient accumulation (§0.9) runs $G$ micro-batches through forward and backward and steps once. When micro-batch $k$ holds $n_k$ counted tokens, the gradient of the big batch is the token-weighted average $\sum_k n_k \nabla\bar{\mathcal{L}}_k / \sum_k n_k$ of the micro-batch mean-loss gradients, which equals the plain mean over $k$ only when every $n_k$ is the same.

The memory cost of this state is worth knowing by heart. Adam keeps two extra numbers per parameter, $\mathbf{m}$ and $\mathbf{v}$ , in float32, so with float32 weights and gradients training costs 16 bytes per parameter before any activations, about 2 GB for GPT-2 small and 112 GB for a 7B model. §10.1 builds the full memory budget from this number.

Here are the schedules of this section's run, of the GPT-2 reproduction, and of two alternatives. Load a preset and read the rate at any step.

The next cell collects the three optimizer pieces in their §0.6 form, compacted, so the rest of the section can use them.

Python · cell 3

A whole training loop, in NumPy

We now have every piece except backpropagation through a Transformer, which autograd will do for us in PyTorch. To see the loop work in the browser, train the simplest language model there is with exactly this machinery, a bigram model whose logits for the next character are a learned row $\mathbf{W}_{x_t,:}$ of a $V \times V$ table, so it predicts from the current character alone (§3.1). Its gradient is the cross-entropy gradient, scatter-added into the rows that were used. We train it three times, with three peak learning rates.

Python · cell 4

Three different stories in one picture. At 0.003 the model is still sliding down at the end, both curves together, which is a learning rate too low for the budget. At 0.3 the training loss drops lowest, but validation (the dashed line) ends worse, because the big steps let the table memorize rare transitions of the seven training sonnets. At 0.03 the validation loss reaches about 2.6 nats, close to the 2.55 that a well-smoothed bigram count model reaches on this split (§6.6 builds it). The model has learned all the one-character statistics there are. To do better it needs context, which is what the Transformer is for.

One optimizer step, assembled

The loop above hides the hardest part of a real step, because it has one batch, no masked targets, and a single parameter tensor. A GPT step has all three complications at once. Here is the order that every correct implementation follows, and the reason for each position in it.

Set the learning rate for this step from the schedule, before anything uses it.
Accumulate. For each of the $G$ micro-batches, run forward and backward and add its gradient, weighted by its share of the counted tokens. In PyTorch you count the targets of all $G$ micro-batches first, then call backward() on each micro-batch's summed token loss divided by that total.
Clip the combined gradient by its global norm, and log the norm measured before clipping. Clipping a single micro-batch, or clipping after the update, protects nothing.
Step AdamW once, which advances its step counter and with it the bias correction.
Zero the gradients, so the next step starts from nothing. Forgetting this silently accumulates every step into the next.

The step also reports two numbers for the log, the training loss as the token-weighted mean over all micro-batches (the last micro-batch's loss alone is a noisy and biased reading), and the pre-clip gradient norm.

Interview

One complete training step

Write train_step(params, opt_state, micro_batches, step, grad_fn, cfg), one full optimizer step of a GPT-style loop, returning (loss, grad_norm, lr).

params is a list of float arrays. Update them in place.
opt_state is a dict {"m": [...], "v": [...], "t": 0} holding Adam's moments (arrays shaped like params, zeros at first) and the number of optimizer steps taken so far. Update it in place.
micro_batches is a list of (x, y) pairs. Targets equal to -100 are ignored, so micro-batches hold different numbers of counted tokens (at least one each).
grad_fn(params, x, y) returns (mean_loss, grads, n), the mean loss over the n counted targets of that micro-batch and its gradients (a list shaped like params), which is what loss.backward() on a mean loss gives you.
cfg has the keys max_lr, min_lr, warmup, total, beta1, beta2, eps, weight_decay, clip.

In order, set the learning rate with lr_at(step, max_lr, min_lr, warmup, total) (provided, the §0.6 schedule), combine the micro-batch gradients into the gradient of the mean loss over all counted tokens, clip that gradient by its global norm (multiply by clip / (norm + 1e-6) when the factor is below 1), and apply AdamW as PyTorch does (add one to t, shrink each tensor with two or more dimensions by 1 - lr * weight_decay, then take the bias-corrected Adam step). Return the token-weighted mean loss, the global norm before clipping, and the learning rate you used.

import numpy as np

def train_step(params, opt_state, micro_batches, step, grad_fn, cfg):
    """returns (token-weighted mean loss, global grad norm before clipping, lr)"""
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

Mixed precision, for a GPT

A GPU runs matrix multiplies in 16 bits many times faster than in float32 (an A100's tensor cores are rated at 312 dense TFLOPS in bfloat16 against 19.5 TFLOPS in plain float32), so every real run uses mixed precision, whose formats and tools you met in §0.9. The rule it follows comes from one experiment that §8.5 runs in full. A bfloat16 weight of 1.0 that receives a thousand updates of 0.001 never moves, because bfloat16 has 8 significant bits and rounds away any relative change below about one part in 256. Late in training, with a decayed learning rate, most updates are that small. So the master weights and Adam's moments stay in float32, the matmuls of the forward and backward passes run in bfloat16 under torch.autocast, and reductions run in float32. float16 has a far narrower range and also needs loss scaling, with the gradients unscaled before clipping.

Two details are specific to a GPT. Upcast the logits to float32 before the cross-entropy, since the log-sum-exp over 50,257 entries is exactly the kind of long reduction that 8 significant bits get wrong (that's the .float() in the model below). And because the logits are the largest tensor in the step, production code fuses the upcast and the loss into one chunked kernel instead of materializing a float32 copy of all of them.

Training our GPT

Everything comes together in PyTorch. First, the sanity check that should precede every real run. Take one batch and train on it alone. A correct model and loss must be able to memorize it, driving the loss toward zero. If it can't, something between the data and the optimizer is broken, and no amount of tuning will fix it (§0.7).

🔥 In PyTorchfor a machine with PyTorch

The loss on the single batch goes from 4.02 (that's $\ln 56$ ) to about 0.01. The plumbing works. Now the real loop, the five steps of the assembled step in PyTorch. Decay groups, AdamW with GPT-3's betas, warmup and cosine decay (reusing lr_at from the NumPy cell), two micro-batches of accumulation, clipping at 1.0, and bfloat16 autocast when a GPU is present. It runs in a couple of seconds on a laptop CPU.

🔥 In PyTorchfor a machine with PyTorch

The training loss falls from 4.06 to about 2.3 and the validation loss to about 2.53. Notice the gap opening between them. With only 4,300 training characters, the GPT soon knows the seven sonnets better than it knows sonnet 116, and its held-out loss barely beats the bigram's 2.55. More steps would widen the gap, not close it. That is the data-limited regime, and the cure is more data, not a bigger model or a longer run. The widgets in §6.4 to §6.6 use a model with this architecture and training code, with a 128-character context and dropout added, trained on all 154 sonnets (about 90,000 characters) for 3,000 steps. It reaches 1.62 nats on ten held-out sonnets.

Reading loss curves

A loss curve answers three questions. Is the model learning (the training loss falls)? Is it learning something general (validation follows training)? And is the optimization healthy (the curve is smooth and its noise is stable)? Each classic failure has a signature shape. The widget below trains a tiny neural bigram model live in your browser, with real data and real AdamW. Pick a mystery curve and diagnose it, or set the knobs yourself.

what you see	likely cause	what to do
loss flat at $\ln V$ from step 0	no learning signal (learning rate 0, frozen parameters, gradients not reaching the weights)	check that gradients are nonzero and the optimizer holds the model's parameters
loss near zero within a few hundred steps, train and val	targets visible to the model (forgotten shift, leaking mask)	check one (input, target) pair and the causal mask (§6.1, §6.2)
slow, almost linear decline	learning rate too low	raise it, the loss should drop faster early
noisy, high plateau, or rises after an early drop	learning rate too high	lower the peak, lengthen warmup
sudden jump that does not recover	loss spike, instability	lower the learning rate, check clipping, restart from a checkpoint before the spike (§8.8)
training falls, validation turns up	overfitting	more data, fewer epochs, regularization, or stop at the best validation step
steps or periodic bumps in the curve	data order effects (unshuffled shards, repeated epochs)	shuffle documents, check the loader's position on resume
validation equal to training, both low	leakage between the splits	split by document and deduplicate (§6.1)

So, the puzzle from the top. A loss that falls, flattens, then jumps to 4 and stays there is a loss spike from which the run did not recover. Notice where it lands. For our 56-token vocabulary, 4 is almost exactly $\ln 56 = 4.03$ , the loss of a uniform guess, so the model has not merely got worse, it has lost everything it learned, typically because one huge update pushed the weights into a region where the logits no longer depend on the input. The learning rate was too high for that stage of training, or one bad batch produced a huge update that the optimizer's moments then amplified. Clipping and a lower peak learning rate make it much rarer. At scale people also restart from a checkpoint before the spike and skip the offending batches.

🎯 Interview questions on this section

Walk me through one optimizer step of a GPT training loop, in order.

Set the learning rate for this step from the warmup-plus-cosine schedule. For each of the $G$ micro-batches, run the forward pass under bfloat16 autocast, compute the cross-entropy over counted targets with the logits upcast to float32, and call backward() on the summed token loss divided by the total counted tokens of all $G$ micro-batches, so the accumulated gradient is the big-batch mean. Under float16, unscale. Clip the global gradient norm at 1.0 and log the norm from before clipping. Take one AdamW step, decaying only the tensors with two or more dimensions. Zero the gradients. Log the token-weighted training loss, the learning rate, the norm, and the tokens seen so far.

Your GPT-2-sized model starts training at a loss of 25. What do you suspect?

A freshly initialized model should start near $\ln V = \ln 50{,}257 = 10.82$ , because small random logits give a nearly uniform prediction. A loss of 25 means the logits are already large, so the model is confidently wrong before learning anything. The usual causes are an initialization far larger than 0.02 (for example an untied head left at a default init, or a missing $1/\sqrt{2L}$ factor in a deep model), a missing final LayerNorm, or a missing $1/\sqrt{d_k}$ in attention. Check the logits' standard deviation at step zero before tuning anything.

Which parameters should get weight decay in a GPT, and why?

Weight matrices and embeddings, the tensors with two or more dimensions. Not biases or LayerNorm gains. A LayerNorm gain's natural value is 1, and decaying it toward 0 fights the normalization, while biases shift activations without adding capacity. Both are a tiny share of the parameters, so they contribute little to overfitting anyway.

How does gradient accumulation work, and when is dividing each micro-batch loss by G wrong?

Run $G$ micro-batches, let gradients sum across backward() calls, and step once. Dividing each micro-batch's mean loss by $G$ reproduces the big-batch gradient only when every micro-batch has the same number of counted tokens. With padding or masked tokens the right gradient is the token-weighted average $\sum_k n_k \nabla\bar{\mathcal{L}}_k / \sum_k n_k$ . In practice you sum per-token losses and divide by the total count over all micro-batches.

Your loss curve shows training loss falling steadily while validation loss starts rising after step 2,000. What's happening and what do you do?

Overfitting, the model is memorizing the training set. Options are more or more diverse data (the real fix in pretraining, where each token is usually seen once), stopping at the best validation checkpoint, more regularization (weight decay, dropout), or a smaller model. Before any of that, check for a split bug, such as duplicated documents, which shows up differently, with validation suspiciously close to training.

The loss spikes from 3.1 to 7 at step 40,000 and comes back down over a few hundred steps. Do you worry?

Yes, even though it recovered. A spike is one or a few enormous updates, usually from a bad batch or from logits and attention scores that have drifted large, and recovery often leaves the run slightly worse than before. Look at the logged gradient norm, which usually climbs before the loss does, and at the data at that step. Common remedies are restarting from a checkpoint before the spike while skipping the offending batches, a lower peak learning rate, and stabilizers such as QK-norm or a z-loss on the logits (§8.8). A spike that lands near $\ln V$ and stays there means the run has lost everything and must be restarted.

How big is the logits tensor for GPT-2 at B=16, T=1024, and why does it matter?

$16 \times 1024 \times 50{,}257 \approx 8.2 \times 10^8$ numbers, 3.3 GB in float32, and the gradient is the same size. It's often the largest activation in the whole step, so implementations compute the loss in chunks or with a fused kernel that never stores the full tensor, and they upcast logits to float32 just for the loss.

🧭 What to remember

The loss is cross-entropy over the vocabulary, averaged over counted targets only, computed with a stable log-sum-exp. Its gradient is softmax minus one-hot, divided by the counted number.
The optimizer pieces come from Chapter 0. GPT-2-style runs use AdamW with $\beta_2 = 0.95$ and $\lambda = 0.1$ on matrices only, warmup then cosine decay, and clipping at 1.0, and keep 16 bytes per parameter of training state.
One correct step sets the learning rate, accumulates token-weighted gradients, clips the global norm (logged before clipping), steps AdamW once, and zeros the gradients.
Train in bfloat16 with float32 master weights, and upcast the logits to float32 for the loss.
Overfit one batch before anything else, then read the train and validation curves for the signatures of each failure.
Next, §6.4 turns the trained model's distributions into text, and shows why the way you sample matters as much as the model.

Loshchilov & Hutter (2019). Decoupled Weight Decay Regularization. AdamW and why L2 and weight decay differ under Adam.
Brown et al. (2020). Language Models are Few-Shot Learners. Appendix B gives the GPT-3 training settings that GPT-2 reproductions reuse.
Micikevicius et al. (2018). Mixed Precision Training. Master weights, loss scaling, and float32 accumulation.
Karpathy. llm.c GPT-2 (124M) reproduction. Every hyperparameter of a real run, with its reason.
Hugging Face (2024). Fixing Gradient Accumulation. The token-count bug in a widely used trainer, and the fix.
Karpathy. A Recipe for Training Neural Networks. The habits behind "overfit one batch first" and reading curves.

Zero to Frontier is made for larger screens