Building a GPT from Scratch

With the architecture built (Chapter 13) and text tokenized (Chapter 14), we can finally train. The pretraining objective for a decoder-only Transformer is the simplest imaginable: given a sequence of tokens, predict the next one at every position. This single objective, scaled up, produces all the capabilities of modern LLMs.

Because of the causal mask (Chapter 12), every position predicts its successor using only earlier tokens, and all T predictions are computed in a single forward pass. This is the efficiency that makes the decoder-only design so attractive: one forward pass yields T training signals, one per token.

During training we feed the true tokens and ask the model to predict each next token — this is teacher forcing. Implementation reduces to a one-position shift: the inputs are tokens x₁...x_{T-1} and the targets are x₂...x_T. The model's logits at position t are scored against the actual token at position t+1.

Transformers are unusually sensitive to the learning rate early in training. Starting at the full learning rate often causes immediate divergence. The fix — learning-rate warmup — ramps the rate up linearly from near zero over the first few thousand steps, then decays it. This single trick is essential for stable Transformer training.

At initialization the model's predictions are random and the gradients are large and noisy. Adam's second-moment estimate (the running variance of gradients) has not yet stabilized, so early updates can be wildly miscalibrated. A large learning rate amplifies these bad early updates, pushing the model into a region from which it cannot recover. Warmup gives Adam's statistics time to settle before applying full-strength updates.

The standard modern schedule combines linear warmup with cosine decay: ramp up linearly to a peak learning rate over the warmup steps, then decay smoothly following a cosine curve down to a small final rate. This schedule trains the model fast once warmed up, then anneals to a low rate for fine convergence.

Chapter 2 derived Adam; here we focus on AdamW, the variant used to train essentially every modern LLM. AdamW's key correction — decoupling weight decay from the gradient update — makes regularization behave correctly with adaptive learning rates.

Plain Adam with L2 regularization adds λθ to the gradient — but then this term gets divided by √v̂ along with everything else, so parameters with large gradients get less regularization. AdamW (Loshchilov & Hutter, 2017) fixes this by applying weight decay directly to the parameters, separately from the adaptive gradient step.

Even with warmup and a good optimizer, training can produce occasional huge gradients — from a difficult batch, a rare token, or accumulated instability. A single enormous update can undo thousands of steps of progress or send the loss to NaN. Gradient clipping caps the gradient's magnitude, keeping every update bounded.

The standard method clips by global norm: compute the L2 norm of all gradients concatenated together, and if it exceeds a threshold, scale every gradient down by the same factor so the total norm equals the threshold. This preserves the gradient's direction while bounding its size.

Chapter 5 introduced the floating-point formats and the mechanics of mixed precision. Here we apply them to Transformer training. The goal: do the bulk of computation in 16-bit precision (bf16 or fp16) for speed and memory, while keeping a master copy of the weights and the optimizer states in 32-bit for accuracy.

For LLM training, bf16 has decisively won. Its full fp32 exponent range means gradients and activations almost never overflow, eliminating the loss-scaling machinery that fp16 requires. The cost — fewer mantissa bits, so less precision — turns out not to matter for training, where range is more important than precision.

Two complementary techniques extend what fits on a given GPU. Gradient checkpointing (introduced in Chapter 11) trades compute for activation memory; gradient accumulation trades steps for a larger effective batch size. Together they let you train models and batch sizes that would otherwise not fit.

To use a large batch size that does not fit in memory, split it into smaller micro-batches, accumulate their gradients, and step the optimizer only after processing all of them. The effective batch size is micro_batch × accumulation_steps, achieving the statistics of a large batch with the memory of a small one.

Batch size and learning rate are not independent. A larger batch gives a less noisy gradient estimate, which permits — and often requires — a larger learning rate. Getting their relationship right is central to efficient training, and the relationship has been studied extensively.

Two heuristics relate learning rate to batch size. The linear scaling rule (Goyal et al., 2017) says lr should scale linearly with batch size. The square-root rule says lr should scale with √(batch size), which better matches the noise structure of SGD. In practice, the square-root rule holds over a wider range for adaptive optimizers.

Note the inverse relationship: larger models use smaller peak learning rates. This reflects the empirical finding (and theoretical work on μP, Chapter 10) that the optimal learning rate decreases as width grows. Hyperparameter transfer techniques aim to tune on a small model and transfer the settings to a large one, saving enormous compute.

At scale, training runs span weeks and cost millions of dollars, and they do not always go smoothly. Loss spikes — sudden jumps in the loss — are common, and full divergence to NaN can destroy a run. Knowing how to diagnose and recover is essential practical knowledge.

The standard recovery for an unrecoverable spike, used in the training of models like PaLM and OPT: roll back to the last good checkpoint, skip the batches that caused the spike, and resume — sometimes with a temporarily lower learning rate. Because the spike is often triggered by specific pathological data, skipping it usually lets training continue cleanly.

A training run is only as good as your ability to see inside it. Experienced practitioners watch a small set of signals that, together, reveal the health of the optimization. Logging these from the start turns debugging from guesswork into diagnosis.

Beyond loss, periodically evaluate on held-out validation data and on downstream benchmarks. Validation perplexity confirms generalization; downstream task metrics (even simple ones) confirm that decreasing loss translates into useful capability. Chapter 21 covers evaluation during pretraining in depth.

We now assemble everything — the objective, schedule, optimizer, clipping, mixed precision, accumulation, and checkpointing — into a single training loop. This is the practical recipe for training a GPT-style model, and it is the loop that, scaled up, trains real LLMs.

Exercises 1–10 are pen-and-paper or derivations; 11–20 require code.

Further reading: “Decoupled Weight Decay Regularization” (Loshchilov & Hutter, 2017) for AdamW. “Attention Is All You Need” (Vaswani et al., 2017) for the original warmup schedule. “An Empirical Model of Large-Batch Training” (McCandlish et al., 2018) for the critical batch size. “OPT: Open Pre-trained Transformer Language Models” (Zhang et al., 2022) and its training logbook for a candid account of real large-model training. Andrej Karpathy's nanoGPT for a complete, readable reference implementation of this chapter's loop.

Next → Chapter 16: Scaling Laws

You can now train a Transformer — but how big should it be, how much data does it need, and how much compute will it take? Chapter 16 answers these with scaling laws: the remarkably precise power-law relationships between model size, dataset size, compute, and loss. We will derive the Chinchilla compute-optimal recipe, understand the 6N rule for training FLOPs, and see how scaling laws turn the art of choosing model size into a predictable science — the quantitative foundation that guides every frontier training run. This chapter closes Part III.