Building a small LLM

Seven chapters in, the pieces are on the bench. Chapters 1-6 built the architecture: embeddings, attention, the transformer block, positional encoding. Chapter 7 built the data pipeline: web text filtered, deduplicated, decontaminated into a clean corpus. What is missing is the part where the architecture learns from the data. That gap is this chapter.

The training loop is short. Maybe sixty lines for a small model. But every line carries decisions that took roughly a decade of LLM research to settle. The loss function, cross-entropy on next-token prediction, is settled. The optimizer, AdamW with specific betas, is settled. The schedule, linear warmup into cosine decay, is settled. Gradient clipping at norm $1.0$ is settled. These choices were not obvious in 2017; by 2024 they are the default that every open frontier lab converges on. This chapter is the walk through why each one matters.

The chapter’s reference implementation is GPT-2-style, a faithful, simplified nanoGPT, that runs on a laptop. After about an hour of training on a small corpus, the model goes from gibberish (“xqz mfn the the the”) to grammatical English. The marquee widget in section 7 lets you scrub through that trajectory: watch the loss drop, watch the generated samples improve, watch random noise turn into a working language model in real time.

The chapter is also the most code-heavy in the tutorial. Six runnable code blocks, each implementing one of the components above. The earlier chapters were “what the model is”; this chapter is “how to make it work.”

The setup: what we’re about to do

What the reader has, after seven chapters, is a parameterized function. A transformer block (Chapter 5) stacked $N$ times, with positional information either added to the input (sinusoidal or learned) or applied inside attention (RoPE, Chapter 6). A tokenizer (Chapter 3) and embeddings (Chapter 2). A training corpus from Chapter 7, billions of tokens of curated web text, ready to feed in. The architecture is wired correctly. A forward pass runs cleanly and produces a probability distribution over the vocabulary at every position. What it does not do, yet, is anything useful.

Every weight matrix in that architecture is still random. The Q, K, V projections inside every attention head; the FFN up- and down-projections; the embedding rows; the layer norm scales. All sampled from a small-variance normal at initialization. The model’s outputs reflect that: feed it a prompt and the next-token distribution is nearly uniform, with maybe a slight bias from however the embedding lookup happened to land. The function has the right shape to be a language model. It does not yet have the right values to be a useful one.

Training is what moves the values. Specifically: training defines a loss that measures how badly the model predicts the corpus, computes the gradient of that loss with respect to every parameter, and takes a small step in the direction that reduces the loss. Repeat for billions of steps. Slowly the random initialization gets sculpted into weights that produce coherent text. That is the whole game. The chapter unpacks what each of those words means in practice.

The pieces decompose cleanly. A loss function measures prediction error per token; cross-entropy on next-token prediction is the canonical choice. An optimizer uses the gradient to update parameters; AdamW is the modern default. A learning rate schedule controls how big each step is at each point in training; linear warmup into cosine decay is the standard recipe. A few stability tricks keep things from blowing up: gradient clipping, mixed precision, gradient accumulation. The full training loop assembles these into the iteration that runs for the whole training budget. Sampling turns the trained model back into text at inference time.

“Small” is worth pinning down. In 2024, frontier models are tens of billions to roughly a trillion parameters; LLaMA-3 8B is 8 billion. This chapter targets ten million to a hundred million, three to four orders of magnitude smaller. Small models train on a laptop in hours; frontier models train on GPU clusters for weeks. The training recipe, though, is the same. The loss is the same loss; the optimizer is the same optimizer; the schedule is the same schedule. What changes at the frontier is the engineering (distributed training, custom kernels, careful failure recovery), covered in Chapters 9 and 10. The mathematics of training are settled, and the small-model recipe in this chapter is the recipe.

TinyStories: how small is small enough

“Ten million to a hundred million parameters” sounds like an odd place to land: small enough to fit on a laptop, but is it small enough to produce anything resembling language? Eldan & Li 2023, arxiv.org/abs/2305.07759, answered this directly. They trained GPT-architecture models with as few as one or two transformer layers and a handful of million parameters, several orders of magnitude below GPT-2, on a purpose-built corpus: TinyStories, a few million short children’s stories generated by GPT-3.5/GPT-4 and restricted to the vocabulary and grammar a three- or four-year-old already understands. The resulting models are fluent. Not GPT-2-fluent, not stylistically rich, but grammatical, coherent across several sentences, and clearly telling a story rather than emitting plausible-looking noise, at a parameter count where training the same architecture on ordinary web text produces gibberish.

The lesson is about capacity matching, not about small models secretly being powerful. A model has to represent the distribution its training corpus draws from, and the complexity of that distribution is set by the data, not by some fixed floor that “language” requires. Full web text spans encyclopedic facts, code, multiple languages, poetry, argument, dialect; fitting it well is what pushes frontier models into the hundred-billion-parameter range. Restrict the distribution to short, simple, repetitive stories built from a small effective vocabulary of concepts and sentence patterns, and a few-million-parameter model has enough capacity to fit it well. TinyStories is the demonstration that “small model” and “incoherent model” are not the same claim; they coincide only when the data is also large and diverse.

This chapter’s own runnable demo, in section 7, makes the identical point at an even smaller scale: a roughly $10$M-parameter, character-level transformer trained on Tiny Shakespeare, about one megabyte of a single author’s stylistic register, well under a hundred unique characters. That constraint, small vocabulary, single style, modest total length, is what lets a model this size visibly learn “language” in a few thousand steps in a browser tab. Word-level modeling over TinyStories’ short children’s stories is the more standard version of the same demonstrator; the recipe, the loss, the optimizer, the schedule, does not change between the two. Only the data and the resulting register of the generated text do.

The first thing training needs is a definition of “wrong.” Section 2 is the loss.

The objective: cross-entropy loss for next-token prediction

The likelihood

A language model is, formally, a probability distribution over sequences. For a sequence of tokens $x_1, x_2, \ldots, x_T$, the model parameters $\theta$ define a joint probability $p_\theta(x_1, \ldots, x_T)$ over the whole sequence. Causal language models factor that joint as a product of per-token conditionals:

$p_\theta(x_1, \ldots, x_T) = \prod_{t=1}^T p_\theta(x_t \mid x_1, \ldots, x_{t-1}).$

Each factor is the model’s predicted distribution for the next token, given all prior tokens. The chain rule of probability makes this factoring exact, no approximation. The architectural choice (causal attention masking, autoregressive generation) is what makes the computation of each factor tractable in a single forward pass.

Training is maximum likelihood estimation on this joint. Given a training corpus, find the $\theta$ that assigns the highest probability to the actual text. Equivalently, and more numerically convenient, minimize the negative log likelihood:

This is one sequence’s contribution to the loss. The full training loss averages over many sequences in the corpus. Minimizing (8.loss) is exactly equivalent to maximizing likelihood; the minus sign and the log are conveniences, not changes of objective.

Cross-entropy per position

At each position $t$, the model produces a vector of logits over the vocabulary, call it $z_t \in \mathbb{R}^V$ where $V$ is the vocabulary size. The softmax of those logits gives the predicted distribution $\hat{p}_t = \softmax(z_t)$. If $y_t$ is the true next token, the loss at position $t$ is $-\log \hat{p}_t[y_t]$, the negative log probability the model assigns to the correct answer.

A useful reframe: $-\log \hat{p}_t[y_t]$ is the model’s surprise at token $y_t$. Low surprise (the model already expected this token) means low loss; high surprise (the model was off-base) means high loss. Training is the program of reducing average surprise across every token in the corpus.

In practice, the loss is computed from logits directly rather than from probabilities. Computing softmax explicitly and then taking the log is numerically dicey: small probabilities round to zero in floating point, and $\log(0)$ is $-\infty$. The fused log-softmax + negative-log-likelihood operation uses the log-sum-exp identity:

$-\log \hat{p}_t[y_t] = -z_t[y_t] + \log\sum_v \exp(z_t[v]),$

which is numerically stable when paired with the standard max-subtraction trick inside the log-sum-exp. Every modern framework (torch.nn.functional.cross_entropy, optax.softmax_cross_entropy_with_integer_labels) provides this fused form. You should always use it.

The compute matters. For a batch of $B$ sequences of length $T$ over a vocabulary of size $V$, the loss involves $B \cdot T$ per-position predictions, each over $V$ classes. For GPT-2 ($V \approx 50{,}000$, $T = 1024$, $B = 32$), that is $32 \cdot 1024 = 32{,}768$ predictions per training step, each a softmax over $50{,}000$ entries. The cross-entropy and the final output projection are a non-trivial fraction of total forward-pass compute, and for very large vocabularies they can dominate. Modern training stacks fuse the projection and the cross-entropy aggressively for this reason.

Why this loss specifically

Cross-entropy is the calibrated loss for probabilistic prediction. Minimizing it is equivalent to minimizing the KL divergence between the model’s predicted distribution and the empirical one-hot distribution that puts all mass on the true next token. The mathematics says: the model is being asked to match a target distribution, not to merely rank the right answer at the top.

The alternatives have problems. MSE on logits treats the output as a regression target, squared error between predicted logits and a target vector. That is the wrong objective: nothing in MSE encodes the constraint that the output should be a probability distribution, and MSE will happily train a model whose logits are correct in scale but uncalibrated in distribution. Hinge loss only penalizes the margin between the correct logit and the next-highest; it produces good classifiers but bad probability estimators. Contrastive losses can train good representations but do not directly optimize next-token likelihood, which is what generation actually needs.

Cross-entropy is the right loss because language modeling is, at the bottom, a calibrated probability prediction problem. Generation samples from the predicted distribution; if the distribution is miscalibrated, sampling is bad. Cross-entropy targets the distribution directly, and the empirical record from 2017 onward shows it works.

The implementation in numpy is short. The computation is exactly the formula: softmax the logits, pick out the probability of the true token at each position, take negative log, average.

The loss is a scalar. To reduce it, training needs the gradient of that scalar with respect to every parameter. That is the forward pass and its backward partner.

The forward pass: tokens to loss

The full forward pass is a sequence of operations that the reader has already met across the previous chapters, now lined up end to end. Token IDs come in; a scalar loss comes out; every intermediate operation contributes a gradient that backprop will compose into a gradient for every parameter.

The pipeline:

1. Token IDs to embeddings (Chapter 2). For each token ID $x_t$, look up its row in the embedding matrix $W_E \in \mathbb{R}^{V \times d}$. Output: a sequence of $d$-dimensional vectors, one per token.
2. Add positional information (Chapter 6). Either add a positional encoding to each embedding (sinusoidal or learned), or, for RoPE, leave the input alone and apply rotation inside attention layers later. The chapter’s reference implementation uses learned positional embeddings, matching the original GPT-2.
3. Run $N$ transformer blocks (Chapter 5). Each block does multi-head attention with residual + layer norm, then a feedforward network with residual + layer norm. The block transforms a sequence of $d$-vectors into another sequence of $d$-vectors of the same shape. Stacking $N$ of them gives the model depth.
4. Final layer norm, applied once after the last block. Standard in Pre-LN architectures; helps stabilize the output projection.
5. Output projection to logits. A linear layer $W_O \in \mathbb{R}^{d \times V}$ maps each $d$-vector to a $V$-dimensional logit vector. Output: a tensor of shape $(B, T, V)$.
6. Cross-entropy loss against the next token. Section 2’s formula, applied at every position. Output: a scalar.

The “input shift” is the small but important detail that ties this together for next-token prediction. The model is asked to predict, at position $t$, the token at position $t+1$. The standard implementation: feed x[:T-1] as input and use x[1:T] as the target. Same sequence, off by one. The model’s logits at position $t$ are scored against the actual token at $t+1$, which is exactly the next-token prediction objective.

Parallel prediction across positions

A single forward pass produces $T$ predictions, not one. The output tensor has shape $(B, T, V)$: at every position $t \in \{0, 1, \ldots, T-1\}$, the model emits a full vocabulary distribution over what token should come next. Position $t$‘s logits are scored against the actual token at position $t+1$, and the cross-entropy loss sums the per-position losses across all $T$ positions and across all $B$ sequences in the batch. A batch of size $B = 32$ at sequence length $T = 1024$ produces $32 \cdot 1024 = 32{,}768$ supervised training examples per step, not $32$. This density is why next-token prediction is so sample-efficient as a self-supervised objective: every token in the corpus is simultaneously an input (it conditions the prediction at every later position) and a label (it is the prediction target for the position before it).

Each of those $T$ predictions sees a different amount of context. Position $0$‘s prediction is computed from just one token of context. Position $1$‘s prediction is computed from two tokens. Position $T - 1$‘s prediction is computed from the full preceding sequence of $T$ tokens. The “how do we reconcile different positions having different context lengths” question dissolves once you see that this is exactly the design: the model is trained simultaneously on prediction problems of every difficulty, from “predict the second token given only the first” (hard, often ambiguous, high loss) to “predict the next token given a thousand tokens of preceding context” (much easier, lower loss). Averaging the loss across all positions teaches the model to be useful at every context length it will face during generation, including the very short prefixes it sees at the start of a generation, where there is little to condition on, and the long prefixes it sees deep into a response.

Causal masking is the mechanism that makes this parallel training honest. Without intervention, the attention mechanism from Chapter 4 would let position $5$ freely attend to position $6$‘s key and value, that is, it would let the model’s prediction at position $5$ peek directly at the very token it is supposed to be predicting. The model would learn a trivial “look one to the right” function and the training signal would be worthless. The causal mask prevents this by zeroing out the attention weights from each position to all later positions. In practice, before the softmax inside the attention block, the model adds a triangular mask that sets the pre-softmax attention scores to $-\infty$ for all pairs $(t, t')$ with $t' > t$; after the softmax, those entries become exactly zero, and so position $t$‘s attention output depends only on positions $0, 1, \ldots, t$. Causal masking is what lets a single forward pass legitimately compute $T$ predictions in parallel: each position is forced to predict from a causally-restricted prefix that exactly matches the situation it would face at inference time.

A useful mental picture: a forward pass is $T$ separate prediction problems running side by side, sharing weights and sharing the computation graph up to the final per-position projection that produces each position’s logits. Causal masking is what severs the lateral information flow that would otherwise let those $T$ problems trivially leak labels into each other. The cost of the masking is essentially nothing, a static triangular mask added before softmax, the same mask reused at every layer and every step of training, and the benefit is the entire reason transformers train so efficiently: $T$-fold supervision from a single forward-backward pass.

The backward pass is the mirror image of the forward. Auto-differentiation (Chapter 1) walks the computation graph backward from the loss scalar, applying the chain rule at each operation. Cross-entropy’s gradient with respect to logits is $\hat{p}_t - \mathbf{1}_{y_t}$, the predicted distribution minus the one-hot true distribution. That tidy expression flows backward through the output projection, through the final layer norm, through the stack of blocks, all the way down to the embedding matrix. At each parameter, the chain rule accumulates the parameter’s contribution to the loss. After backward, every weight in the model has an attached gradient and the optimizer can step.

Memory is the constraint that dominates this whole computation. For backprop to work, the framework must keep the activations produced during forward, every block’s outputs, every attention’s keys and values, every intermediate, so they are available when their gradients are needed. The memory for those activations grows linearly with depth, linearly with sequence length, and roughly linearly with batch size. For a 100M-parameter model at sequence length $1024$ and batch size $32$, the activations alone can run into the gigabytes. This is why gradient checkpointing (recompute activations on the backward instead of storing them) and mixed precision (store activations in 16-bit instead of 32-bit) matter at scale. Section 6 covers mixed precision; Chapter 10 covers checkpointing and the rest of the memory tooling.

That is the forward pass at the level of operations. The full implementation is the capstone code in section 7; it composes the building blocks from prior chapters with the loss from section 2 and runs them end to end. With the loss computable and gradients available, the next question is how to take steps.

The optimizer: from SGD to Adam to AdamW

Why not SGD

Stochastic gradient descent is the textbook starting point. At each step, with gradient $g_t = \nabla_\theta \mathcal{L}(\theta_t)$:

$\theta_{t+1} = \theta_t - \alpha g_t.$

Subtract the gradient, scaled by a learning rate. Simple, universal, theoretically well-understood. SGD works fine for many problems, image classification with CNNs, for instance, was trained with SGD + momentum for years. It does not work well for transformers, and the reasons are specific.

Adaptive scale is missing. SGD applies the same learning rate to every parameter. Inside a transformer, different parameters need wildly different update magnitudes: the embedding rows for rare tokens see almost no gradient signal and need large updates when they do; the layer norm scales get strong signal every step and want small updates. A single learning rate that works for both is hard to find, and tuning the rate becomes prohibitively expensive at scale.

No momentum. Vanilla SGD has no memory of prior gradients. Each step is determined by the current minibatch alone. For transformers, gradients across minibatches are noisy: different sequences in different batches give wildly different per-parameter signals. Plain SGD bounces around in response, and convergence suffers. Momentum-SGD ($m_t = \beta m_{t-1} + g_t$, $\theta_{t+1} = \theta_t - \alpha m_t$) partly fixes this by averaging gradients across steps, but it still leaves the adaptive-scale problem.

Tuning sensitivity at scale. Even with momentum, SGD on a transformer requires careful per-layer or per-parameter learning rate scheduling. Practitioners report that getting SGD to train a transformer to GPT-2-class loss values requires substantial tuning that adaptive optimizers do for free. Past a few hundred million parameters, almost no one tries.

The modern transformer recipe is Adam (Kingma & Ba 2014, arxiv.org/abs/1412.6980) and its descendant AdamW (Loshchilov & Hutter 2017, arxiv.org/abs/1711.05101). Both add per-parameter adaptive learning rates on top of momentum, and both are now overwhelmingly the default.

Adam

Adam maintains, for every parameter, two extra state variables: the first moment $m_t$ (an exponentially weighted moving average of the gradient, the momentum-like term) and the second moment $v_t$ (an exponentially weighted moving average of the squared gradient, the adaptive-rate term).

The updates:

$m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t,$ $v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2.$

At step $t$, $m_t$ is a smoothed estimate of the recent gradient direction; $v_t$ is a smoothed estimate of the recent gradient magnitude (squared). Both are initialized to zero, which biases the early estimates toward zero (an EMA initialized at zero needs many steps to “warm up” to the true value). Adam corrects this with bias-corrected versions:

$\hat{m}_t = \frac{m_t}{1 - \beta_1^t}, \qquad \hat{v}_t = \frac{v_t}{1 - \beta_2^t}.$

The parameter update uses the bias-corrected estimates:

$\theta_{t+1} = \theta_t - \alpha \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}.$

Two properties of this formula are doing the work. Direction smoothing: $\hat{m}_t$ averages over recent gradients, damping the noise that single-batch SGD would propagate. Per-parameter adaptive scale: $\sqrt{\hat{v}_t}$ normalizes by historical gradient magnitude, so a parameter with consistently large gradients gets smaller effective steps and a parameter with consistently rare gradients gets larger effective steps. The result is that Adam needs much less learning-rate tuning than SGD: the optimizer figures out the per-parameter scale on its own.

The Adam defaults from the original paper are $\beta_1 = 0.9$, $\beta_2 = 0.999$, $\epsilon = 10^{-8}$. Transformer practice has shifted: $\beta_2 = 0.95$ is now standard for LLMs, on the empirical observation that the lower $\beta_2$ adapts faster to gradient magnitude changes and produces better-trained models.

AdamW: decoupled weight decay

Adam plus L2 regularization should, in principle, be a small modification: add $\lambda \theta$ to the gradient before the optimizer step, and the optimizer behaves as if minimizing $\mathcal{L} + \frac{\lambda}{2} \|\theta\|^2$. This is exactly what most implementations did from 2014 to 2017, and it is exactly what AdamW (Loshchilov & Hutter 2017) showed is wrong.

The problem is subtle. The L2 gradient $\lambda \theta$ gets added to the data gradient $g_t$, the sum gets squared and accumulated in $v_t$, and then the update is divided by $\sqrt{\hat{v}_t}$. The L2 contribution gets divided by the same $\sqrt{\hat{v}_t}$ that the data gradient is divided by, which means the effective weight decay rate becomes parameter-specific, large for parameters with small $\hat{v}_t$, small for parameters with large $\hat{v}_t$. That is not what L2 regularization is supposed to do. L2 regularization is supposed to apply a uniform shrinkage; Adam + L2 applies a non-uniform shrinkage shaped by gradient history.

AdamW’s fix is to apply weight decay as a separate, decoupled step:

$\theta_{t+1} = \theta_t - \alpha \left(\frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda \theta_t\right),$

or equivalently:

$\theta_{t+1} = (1 - \alpha \lambda) \theta_t - \alpha \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}.$

The first term is the decoupled decay: shrink $\theta$ by a constant factor $(1 - \alpha \lambda)$ at every step, independently of Adam’s adaptive scaling. The second term is the standard Adam update, without the L2 contribution folded in. The decay is now uniform across parameters and reflects the textbook intent of weight decay as a regularizer.

Empirically AdamW produces measurably better-trained models than Adam + L2 at the same nominal hyperparameters. The Loshchilov-Hutter paper showed the gap clearly on image classification; subsequent transformer practice (GPT-3, LLaMA, Mistral, every open frontier model) confirmed it for language modeling. The convention has flipped: AdamW is the default, Adam-plus-L2 is the historical mistake, and weight decay rates for LLMs are now meaningfully higher ($\lambda = 0.1$ is standard) than they were in pre-2018 practice.

The numpy AdamW is fifteen lines. Two state buffers per parameter, the EMAs, the bias correction, the decoupled decay step.

The learning rate $\alpha$ in AdamW is not a constant. It changes at every step, following a fixed schedule. That schedule is the next section.

Learning rate schedules: warmup + cosine decay

The peak learning rate is a hyperparameter, and it depends on model size. Empirically, $\alpha_{\max} \approx 6 \cdot 10^{-4}$ works for $\sim 125$M-parameter models; $\alpha_{\max} \approx 3 \cdot 10^{-4}$ for $\sim 1.3$B; $\alpha_{\max} \approx 1 \cdot 10^{-4}$ for $7$B and above. The pattern is that larger models want smaller peak rates, which holds across most published recipes from GPT-2 onward. The chapter’s reference model is small (around 30M parameters), and the configuration uses $\alpha_{\max} = 6 \cdot 10^{-4}$.

The schedule has two phases. A linear warmup ramps the learning rate from zero to $\alpha_{\max}$ over the first $W$ steps. A cosine decay then smoothly drops the rate from $\alpha_{\max}$ down to a floor $\alpha_{\min}$ over the remaining $S - W$ steps, where $S$ is the total training budget in steps:

Warmup. At step zero, every parameter in the model is random. The gradient at step zero is, correspondingly, garbage, high variance, partly because the loss landscape near random initialization has steep features, partly because batch statistics dominate when no parameter has been refined yet. Taking a full-size optimizer step on that gradient frequently destabilizes training (loss spikes, NaN propagation, parameter blowups). Linear warmup starts at $\alpha = 0$, so the first step takes essentially no step, and the rate grows linearly until step $W$ when the model has “found a reasonable region” of parameter space and the full peak rate is safe. Warmup was originally introduced (Goyal et al. 2017, arxiv.org/abs/1706.02677) for large-batch training, where the problem is most acute; the lesson generalized and is now standard at every batch size.

Cosine decay. After warmup, the rate decays following a half-cosine: a smooth, monotonic curve from $\alpha_{\max}$ down to $\alpha_{\min}$. The intuition: late in training, the model is close to a good minimum and wants to fine-tune with small steps; cosine provides that without the abrupt transitions of step-decay schedules (which have been shown to cause loss “step changes” right at the boundary). Cosine annealing was introduced by Loshchilov & Hutter 2016 (arxiv.org/abs/1608.03983) and has been the default transformer schedule since.

Typical values. Warmup of $W = 2000$ steps for medium-scale models is standard; very large models use up to $10{,}000$. The floor $\alpha_{\min}$ is conventionally about $10\%$ of $\alpha_{\max}$, small enough to fine-tune, large enough to keep making progress. The total steps $S$ is the training budget, chosen to consume the available training tokens at the chosen batch size. For a $50$M model trained on $1$B tokens with $32$k tokens per step, $S = 1$B$/32$k $\approx 31{,}000$ steps; for the chapter’s demo configuration ($S = 100{,}000$) that would be a $3.2$B-token budget.

The schedule controls step size. There is also the question of what happens when a step would be too large in a different sense, when an outlier batch produces a gradient much bigger than usual. That is the next section’s stability story.

Training stability: clipping, mixed precision, accumulation

Three stability techniques are non-negotiable in modern transformer training. None of them changes the loss or the optimizer; each one keeps a particular failure mode from killing the run.

Gradient clipping

Outlier batches happen. Some particular minibatch, maybe one with an unusually long sequence, maybe one with unusual content, produces a gradient that is much larger than the typical gradient. Without intervention, that gradient triggers a much larger parameter update, which in turn destabilizes the model: the loss spikes upward at the next step, sometimes recovers, sometimes propagates into NaN or Inf values and the run is dead.

The fix is to bound the gradient norm. Before the optimizer step, compute $\|g\|_2$, the L2 norm of the full gradient (concatenated across all parameters). If it exceeds a threshold $c$, scale the entire gradient down to norm $c$:

$g \leftarrow g \cdot \min\!\left(1, \frac{c}{\|g\|_2}\right).$

The direction is preserved; the magnitude is capped. Standard $c = 1.0$ for transformers. Smaller values ($c = 0.5$ or $0.1$) are sometimes used for very large models that show occasional gradient spikes; the cost is more conservative updates throughout.

Mixed precision

Modern hardware (NVIDIA Tensor Cores from Volta onward, TPUs, every recent GPU generation) executes FP16 or BF16 matrix multiplications much faster than FP32. Mixed-precision training (Micikevicius et al. 2018, arxiv.org/abs/1710.03740) uses this: forward and backward passes run in FP16 or BF16, which halves memory bandwidth and roughly doubles matmul throughput. The optimizer state, the master copies of the parameters, the AdamW first and second moments, stays in FP32. The hardware does the matmuls in low precision; the optimizer does the updates in high precision. The result: faster training, less memory, no accuracy loss versus pure FP32.

The non-negotiable detail is the FP32 master weights. Pure FP16 throughout loses too much precision in the optimizer state: the second moment $v_t$ accumulates products of small numbers, and FP16’s narrow representable range makes that accumulation drift. Training silently degrades: the loss curve looks fine on the first pass, but the trained model is worse than a pure-FP32 training run would have produced. The mix is what makes mixed precision work; FP16 alone does not.

BF16 (bfloat16) has the same number of exponent bits as FP32, just fewer mantissa bits. It has worse precision than FP16 in the typical range but a much wider representable range. For LLM training, BF16 has effectively replaced FP16 because the dynamic range matters more than precision: gradient values vary across many orders of magnitude, and FP16 sometimes underflows or overflows in a way that BF16 does not. On A100, H100, TPU v4+, BF16 is the modern default.

FP16 training needs a workaround for that narrow range: loss scaling. Multiply the loss by a large constant (say $2^{16}$) before calling backward, so gradients that would otherwise underflow to zero in FP16’s limited exponent range get shifted up into representable territory; unscale the gradients by the same constant right before the optimizer step, so the update itself is unaffected. Frameworks typically make this dynamic, shrinking the scale factor whenever an overflow (an inf or nan gradient) is detected, and growing it again once training has run stably for a while. Loss scaling is a real technique with real engineering, and it is also a symptom of FP16’s fundamental problem: BF16’s exponent range matches FP32’s, so gradients rarely underflow or overflow in the first place, and BF16 training typically skips loss scaling entirely. It is one more reason BF16 has displaced FP16 as the default.

Gradient accumulation

The desired effective batch size is sometimes larger than what fits in GPU memory. A model that can only run forward + backward at batch size $4$ but wants effective batch size $32$ accumulates gradients across $8$ micro-batches before stepping the optimizer:

optimizer.zero_grad()
for accum_step in range(accum_steps):           # accum_steps = 8
    x, y = next(loader)
    loss = compute_loss(model, x, y)
    (loss / accum_steps).backward()             # accumulate scaled gradients
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()

Dividing the loss by accum_steps is the part that occasionally trips people up. The backward pass accumulates gradients into .grad automatically, calling .backward() eight times sums eight gradients. To match the gradient that a single batch of size $32$ would have produced, each of the eight micro-batch losses needs to contribute $1/8$ of the gradient. Hence the division.

Gradient accumulation has no downsides at the loss level (the trained model is identical to one trained at the same effective batch size on bigger hardware), but it has compute downsides: $8$ micro-batches take more wall-clock time than one batch of equivalent size on hardware that can hold it. The technique exists because memory is the binding constraint more often than compute.

The loss, the optimizer, the schedule, and the stability tools are all the pieces of the training loop.

The complete training loop

Putting it together

The training loop is shorter than the prose explaining it. At each step, the loop:

1. Get a batch. The data loader produces input tokens x and target tokens y (shape (B, T) each, y is x shifted by one position).
2. Update the learning rate according to the schedule.
3. Forward pass. Run the model on x, get logits of shape (B, T, V).
4. Compute loss. Cross-entropy of logits against y, averaged over all $B \cdot T$ positions.
5. Backward pass. loss.backward() populates .grad on every parameter.
6. Clip gradients to the configured max norm.
7. Optimizer step. AdamW applies the update.
8. Log periodically. Step number, learning rate, training loss; every few thousand steps, validation loss.
9. Repeat for the configured number of total steps.

That is the entire control flow. A “step” is one optimizer update, one pass through the nine items above. The total number of training tokens is B * T * total_steps. For GPT-2 small ($B = 32$, $T = 1024$, $\text{total\_steps} = 100{,}000$), the loop processes roughly $3.3$B tokens total, about $26$ tokens per parameter at GPT-2’s $124$M, which is in the Chinchilla scaling regime (Chapter 9 covers the derivation).

Batch construction: packing vs padding

Step 1 above, “get a batch,” hides a real design decision. Chapter 7’s corpus is a pile of documents of wildly different lengths, a two-sentence snippet next to a ten-thousand-token article, but the loop needs fixed-shape tensors, (B, T) for x and y, at every step. Two ways to get there.

Padding treats each document as its own sequence. Pad every sequence in the batch up to the longest one (or to a fixed T) with a reserved PAD token, and mask out the padded positions so they contribute nothing to the loss or, inside attention, to what other positions can see. This is correct and simple, and it is standard for fine-tuning and inference batches, where the sequences genuinely are distinct examples (a batch of chat turns, a batch of documents to classify) and the compute wasted on padding is a small fraction of the batch.

Packing is the pretraining default. Concatenate the entire tokenized corpus into one long token stream, insert a separator token (often the same token used for end-of-sequence) between documents, and slice the stream into contiguous windows of length T. Every position in every window is a real token from real text, nothing is padding, so nothing is wasted. The cost is that trained-on windows do not respect document boundaries: a window can start mid-document, end mid-document, or contain the tail of one document immediately followed by the start of the next. Because attention is causal, not bidirectional, the model can technically attend back across that seam, using the unrelated previous document while predicting the first few tokens of the new one, a small amount of cross-document leakage. Production pretraining runs generally accept this rather than pay padding’s compute tax; some reset attention at document boundaries to remove the leakage entirely, at the cost of a more involved data loader.

The chapter’s get_batch() in the exercises is a simplified version of the packing case: data is one long tensor, and every window data[i:i+SEQ_LEN] is a slice of it, with no padding token anywhere. That is the right default to reach for; reach for padding only when the batch’s sequences are genuinely separate examples rather than windows into one continuous stream. One full pass through the concatenated stream is an epoch. Small-model training of the kind this chapter targets often reuses the corpus for several epochs; frontier pretraining runs are usually well under a single epoch over corpora of trillions of tokens, more data than the model will ever see twice.

Running the loop

A trained reader has, after reading the prose so far, an architecture from Chapter 5 and a tokenized corpus from Chapter 7. With those two and a TrainConfig, the reference implementation runs end to end. The code below is intentionally close to nanoGPT, minimal, runnable, and structured to make every step traceable to the prose.

q!Ck;5Wj?n.zUM/x bPL3 GcRO'jY;w.zE:dQs8ux
rYjBkP9q!hZ x?bMNuTL,e 
Ns;3'm.Q?GZ.,LpDxnEMHKj

A few details in the code are worth pointing at explicitly. The parameter-group split (decay and no_decay) is the convention that weight decay applies only to matrices (p.dim() >= 2) and never to biases or layer-norm scales. Biases and norm scales are one-dimensional; decaying them tends to hurt performance because they are meant to take arbitrary task-driven values. Splitting param groups is one line of bookkeeping in PyTorch and worth getting right. The betas (0.9, 0.95) reflect the LLM convention (not Adam’s $0.999$ default). The input shift is implicit, train_data.get_batch() is expected to return (x, y) already paired, with y being x shifted by one position. The logits reshape logits.view(-1, V) flattens the batch and time dimensions into one, F.cross_entropy expects a flattened (B*T, V) input and a flattened (B*T,) target.

Reading the loss curve

The loop prints a number every log_interval steps. Reading that number well, recognizing what a healthy trajectory looks like versus what a broken one looks like, is as important a skill as writing the loop itself.

A healthy run has a recognizable shape. It starts near the random baseline ($\log V$, Section 2), drops steeply over the first few hundred to few thousand steps while the model picks up the easy structure (token frequencies, common bigrams, basic syntax), then flattens into a slower, longer decline as the remaining loss reduction comes from harder, longer-range patterns. Warmup shows up as a small wrinkle right at its own boundary: because the learning rate is still ramping through the first $W$ steps, a run with too short a warmup sometimes shows a visible bump, loss ticks up for a handful of steps right as the rate nears its peak, before resuming its decline once the model adjusts. A plateau, an extended flat stretch where loss stops visibly moving for thousands of steps, is not automatically a problem: late in a cosine schedule, with the rate near $\alpha_{\min}$, slow improvement is expected and the curve legitimately flattens. A plateau early in training, while the rate is still near peak, is more concerning; it often means the learning rate is miscalibrated for the data, or the model has hit a genuine optimization difficulty.

Divergence is the pattern that demands action. Loss jumps, sometimes 2-3x, sometimes to nan, over one or a handful of steps, usually triggered by an outlier batch producing an unusually large gradient. Section 6’s clipping is the first line of defense, but clipping only bounds the damage from a single step; it does not prevent a bad batch from nudging the model into a worse region that then has to be recovered from. Small nanoGPT-scale runs rarely show this pattern; it is far more common at frontier scale, where a run spans weeks across thousands of GPUs and a bad batch, a hardware glitch, or an unlucky data shard becomes a near-certainty over that many steps. Teams training frontier models watch for spikes in close to real time and intervene: skip the offending batch, drop the learning rate, or restart from the last stable checkpoint, a saved copy of the model weights and the full optimizer state ($m_t$, $v_t$, the step counter, not just $\theta$) written to disk periodically during training. Restoring from a checkpoint after a divergence is only faithful if the optimizer state comes back with it: restarting from the saved weights alone but with fresh, zeroed $m_t$ and $v_t$ throws away Adam’s accumulated momentum and second-moment estimates, and the resumed run behaves differently, usually worse, briefly, than if it had never stopped.

The trained model has a forward pass and a probability distribution over the vocabulary at every position. Generating text from it is one more set of choices, about how to convert that distribution into actual sampled tokens.

Generation: sampling from the trained model

The autoregressive inference loop

Generation runs the trained model in a loop, emitting one token at a time. The mechanics differ from training in instructive ways. At training time, a forward pass on a length-$T$ sequence produces $T$ supervised predictions and the loss is computed at every position. At inference time the forward pass produces the same $(B, T, V)$ logits tensor, but only one slice of it matters: the logits at the last position. That row is the model’s prediction for what token should come after the current sequence, and it is what the sampler will turn into the next emitted token. The other $T - 1$ rows of logits get computed and then discarded.

The iteration is: prompt $\to$ forward pass on the prompt $\to$ take logits at the last position $\to$ sample one new token $\to$ append it to the sequence $\to$ run the forward pass again on the now-one-longer sequence $\to$ take logits at the new last position $\to$ sample $\to$ append. Repeat until a stop token or a length budget is hit. After $K$ generated tokens, the model has been run $K$ times, each pass on a progressively longer input. Greedy decoding (next), temperature, and top-$p$ all live inside the “sample one new token” step; the loop around them is the same.

The naive cost is dramatic. Each iteration runs the entire model on the entire sequence so far, every transformer block, every attention head, every FFN, all the way up to the output projection, and uses only the last position’s output. The first $P$ positions in the $(k+1)$-th step did effectively the same computation they did on the $k$-th step; almost all of it is redundant work. For a prompt of length $P$ and $K$ generated tokens, the naive total cost scales roughly as $\sum_{k=0}^{K-1}(P + k) \approx K \cdot (P + K/2)$ forward-pass-tokens, quadratic in $K$ and dominated by repeated re-derivation of identical hidden states from identical prefixes.

The fix is KV caching, covered in Chapter 17. The intuition: attention’s keys and values for past tokens do not change once they are computed, they depend only on prior tokens, and prior tokens are by definition no longer changing, so a serving stack stores those keys and values and only computes fresh keys, queries, and values for the new token at each step. With a KV cache, each decode step does work proportional to one token’s contribution to attention against the stored cache, plus a single-token-wide FFN and projection, rather than re-running the whole model on the whole sequence. The asymptotic improvement is roughly the difference between $O(K^2)$ and $O(K)$ total work for $K$ generated tokens, and the wall-clock factor is comfortably $100\times$ to $1000\times$ at long context. Chapter 17 walks through the implementation, the memory cost (the cache itself is huge at long context), and how it shapes everything downstream, batching, paged memory, speculative decoding.

For this chapter, the takeaway is: yes, the naive generation loop in the next subsection runs the full model on the full prefix at every step. That is what the math says to do, and the small-scale demo code in this chapter does exactly that, it is correct, it is slow, and it is the right starting point. Real serving stacks do something much smarter, and that story is Chapter 17.

Greedy decoding

The simplest decoding strategy is to take the model’s most likely next token at every position. Run the forward pass, pick $\argmax$, append, repeat:

for _ in range(max_new_tokens):
    logits = model(tokens)[:, -1, :]                # (B, V)
    next_token = logits.argmax(dim=-1, keepdim=True)
    tokens = torch.cat([tokens, next_token], dim=-1)

Greedy is deterministic: same prompt, same output every time. It is also boring and frequently degenerate. The model’s argmax is often a high-probability common token; the next argmax is often the same token; after a few steps the output collapses into repetitive loops (“The cat sat on the mat. The cat sat on the mat. The cat sat…”) because greedy decoding finds the model’s preferred filler and gets stuck there. Greedy is rarely the right choice for open-ended generation.

Temperature

Temperature reshapes the distribution before sampling. Before applying softmax, divide the logits by a temperature $T$:

$\hat{p}_t[v] = \frac{\exp(z_t[v] / T)}{\sum_{v'} \exp(z_t[v'] / T)}.$

At $T = 1$, the distribution is unchanged. At $T < 1$, the distribution sharpens, high-probability tokens get even more probability, low-probability tokens lose theirs. At $T \to 0$, the distribution collapses onto the argmax (greedy decoding is the limit). At $T > 1$, the distribution flattens, high and low probabilities move toward uniform. At $T \to \infty$, the distribution becomes uniform.

Typical generation temperatures sit between $0.7$ and $1.0$ for general text. Lower temperatures (around $0.2$ to $0.5$) are used for code generation, where there is usually a single correct continuation and randomness hurts. Higher temperatures ($1.0$ to $1.2$) are sometimes used for creative writing, where surprise is a feature.

Top-k and top-p (nucleus) sampling

Temperature alone is not enough, even at temperature $1$, sampling from the full vocabulary distribution occasionally produces very low-probability tokens that derail the generation. Two filtering schemes address this.

Top-k keeps only the $k$ most likely tokens at each step, masks the rest to probability zero, renormalizes, and samples from the resulting distribution. The hyperparameter $k$ is typically $40$ or $50$. Top-k caps the tail at a fixed size regardless of how peaked or flat the underlying distribution is.

Top-p (nucleus) sampling keeps the smallest set of tokens whose cumulative probability exceeds $p$, masks the rest, renormalizes, and samples from the nucleus. The hyperparameter $p$ is typically $0.9$ or $0.95$. Top-p is more adaptive than top-k: when the distribution is sharply peaked (the model is confident), the nucleus is small and sampling is nearly deterministic; when the distribution is flat (many plausible continuations), the nucleus is large and sampling explores more. The adaptivity matches what people usually want from a sampler, and top-p has largely replaced top-k as the default in modern decoding stacks.

That covers the inference side. The model that section 7’s loop produced now turns into text. The chapter has, in eight sections, walked from a random-initialized transformer to a generating language model. The remaining question is what to do at larger scale.

What we’ve built: and what’s next

Eight chapters in, the parts add up. Chapters 1-6 specify the architecture: embeddings, attention, the transformer block, positional encoding. Chapter 7 specifies the data: filtered, deduplicated, decontaminated web text. This chapter specifies the training: cross-entropy on next-token prediction, AdamW with $(\beta_1, \beta_2) = (0.9, 0.95)$ and weight decay $0.1$, linear warmup into cosine decay, gradient clipping at norm $1.0$, mixed precision for memory and speed, and a sixty-line training loop that runs all of it. Architecture plus data plus training equals a working, small, language model.

“Small” is doing real work in that sentence. The chapter’s reference is $10$M to $100$M parameters, trained on $1$GB to maybe $10$GB of text, runnable on a single GPU in hours. A frontier model in 2024 is hundreds of billions of parameters, trained on tens of trillions of tokens, on a cluster of thousands of GPUs for weeks. The numerical gap is enormous; the recipe gap is small. The same cross-entropy. The same AdamW with $\beta_2 = 0.95$. The same warmup-plus-cosine schedule. The same clip at $1.0$. What changes at scale is the engineering: how to shard a model across many GPUs without bottlenecking on network communication, how to write the matmul kernels that the architecture leans on, how to recover from inevitable hardware failures over week-long runs. The mathematics of training does not change.

The next two chapters cover that engineering. Chapter 9 covers scaling laws (how to choose model size and dataset size given a compute budget, the Chinchilla derivation, the empirical scaling exponents) and distributed training (how to shard parameters, activations, and optimizer state across many GPUs without losing the recipe). Chapter 10 covers training infrastructure (GPU clusters, custom CUDA and Triton kernels for the bottleneck operations, the systems engineering that frontier-scale training rests on). Both chapters build on this one; they take the recipe as given and scale it.

Exercises

Four exercises that build on the chapter’s machinery. Each is a self-contained problem with a starting template; the runnable code lets you iterate without leaving the page. Hints are collapsed; try the problem first.

Exercise 1 (easy): Cross-entropy from logits

Implement cross-entropy loss for next-token prediction. Verify it matches the theoretical random baseline ($-\log(1/V)$ for a uniform distribution over $V$ classes) when the logits are random and the targets are random.

def cross_entropy(logits, targets):
    logits_shifted = logits - np.max(logits, axis=-1, keepdims=True)
    log_sum_exp = np.log(np.sum(np.exp(logits_shifted), axis=-1, keepdims=True))
    log_probs = logits_shifted - log_sum_exp
    n = logits.shape[0]
    return -np.mean(log_probs[np.arange(n), targets])

Exercise 2 (medium): Implement AdamW from scratch

Implement AdamW with bias correction. Verify your implementation converges on a simple quadratic, and observe that weight decay biases the convergence point toward the origin.

class AdamW:
    def __init__(self, params_shape, lr=0.01, beta1=0.9, beta2=0.95, eps=1e-8, weight_decay=0.0):
        self.lr = lr
        self.beta1 = beta1
        self.beta2 = beta2
        self.eps = eps
        self.weight_decay = weight_decay
        self.m = np.zeros(params_shape)
        self.v = np.zeros(params_shape)
        self.t = 0

    def step(self, params, grads):
        self.t += 1
        self.m = self.beta1 * self.m + (1 - self.beta1) * grads
        self.v = self.beta2 * self.v + (1 - self.beta2) * grads**2
        m_hat = self.m / (1 - self.beta1**self.t)
        v_hat = self.v / (1 - self.beta2**self.t)
        params[:] = (1 - self.lr * self.weight_decay) * params - self.lr * m_hat / (np.sqrt(v_hat) + self.eps)

Exercise 3 (medium): Warmup + cosine schedule

Implement the warmup + cosine decay LR schedule. Verify it ramps to peak at the warmup boundary, decays smoothly, and reaches the minimum floor at total_steps.

def lr_schedule(step, max_lr=6e-4, min_lr=6e-5, warmup_steps=2000, total_steps=100000):
    if step < warmup_steps:
        return max_lr * (step / warmup_steps)
    elif step <= total_steps:
        progress = (step - warmup_steps) / (total_steps - warmup_steps)
        return min_lr + 0.5 * (max_lr - min_lr) * (1 + np.cos(np.pi * progress))
    else:
        return min_lr

Exercise 4 (hard): Tiny end-to-end training loop

Train a very small character-level transformer on a tiny corpus end-to-end. The scaffolding (data, model class, hyperparameters) is provided; you implement the training loop and verify the loss drops.

for step in range(TOTAL):
    lr = lr_at(step)
    for g in optimizer.param_groups:
        g['lr'] = lr
    x, y = get_batch()
    logits = model(x)
    loss = F.cross_entropy(logits.view(-1, vocab_size), y.view(-1))
    optimizer.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()
    losses.append(loss.item())
    if step % 50 == 0:
        print(f"step {step:4d}  lr={lr:.2e}  loss={loss.item():.3f}")

What is also worth doing right now, before reading further, is training a model. The reference code in section 7 runs on Tiny Shakespeare ($\sim 1$MB of text, character-level vocabulary, $\sim 10$M parameters) in a couple of hours on a single modern GPU and produces grammatical English. The loss curve drops. The samples improve. The model goes from “fz qm xj” to “Be not afraid of greatness; some are born great.” The pipeline works. Everything from here is making it bigger.