Training

At the end of the last chapter Sprout was full size and completely empty: 17.31 million random numbers, surprised by every token as if it were choosing among all 8,192 at random. In this chapter we fill those numbers in. The whole of training is one short loop: show the model a piece of text, measure how surprised it was, push every number a little in the direction that would have made it less surprised. Repeat 10,070 times.

Before we take the loop apart, let's look at what it does. During training we stopped every so often and asked Sprout to continue four prompts. Here is what it wrote.

Watching Sprout learn

The time-lapse goes through recognisable stages. At step zero the prompt is followed by a jumble of random tokens: the model has no preferences at all. After a couple of dozen steps the most frequent tokens take over: full stops, commas, " the", " was", " and". The model has discovered which tokens are common, the first thing anyone notices about a language. Around step 50 the first templates appear, "Once upon a time, there was a girl named Lily", with quotes in the right places but the grammar still falling apart. By step 100 or so most sentences are grammatical and the dialogue has the right Name: line format, but the story jumps from one thing to another. After a few hundred steps stories stay on topic for a whole paragraph. Everything after that is the slow part: fewer contradictions, longer coherence, better choice of words.

Watch the numbers above the curve as well. "Choices per token" is $e^{\text{loss}}$, the perplexity from chapter 2: the number of equally likely options the model is effectively choosing among. At the start it is in the thousands, even a little above 8,192: a random model is slightly worse than an honest uniform guess. By the end of training it is a handful. That shrinking number is the whole story of this chapter.

One step of training

Here is the heart of train.py, the part that runs 10,070 times. Everything else in the file (arguments, logging, checkpoints) is scaffolding around these thirteen lines:

T = context_at(step, steps, args.context, args.seq_warmup) x, y = train.get(args.batch_tokens // T, T) with ctx: _, loss = train_model(x, y) loss.backward() norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) f = lr_factor(step, steps, args.warmup, args.cooldown) for opt in opts: for g in opt.param_groups: g['lr'] = g['base_lr'] * f opt.step() opt.zero_grad(set_to_none=True) seen += x.numel()

Piece by piece:

1. Window length. context_at says how long the windows are at this step: 128, 256 or 512 tokens. We'll come back to why it changes.
2. A batch. train.get cuts 32768 // T random windows out of train.bin: 256 windows of 128 tokens, or 64 windows of 512. Each window comes with its copy shifted by one token, the answers y. The batch always holds 32,768 tokens, so every step gives the model 32,768 next-token questions at once.
3. Forward pass. The model reads x and returns the loss: the average surprise over all positions of all windows. The with ctx line runs it in bf16 (more on that below).
4. Backward pass. loss.backward() is the backpropagation of chapter 4, applied to 17 million parameters: afterwards every parameter has a .grad, the slope of the loss with respect to it.
5. Clipping. If the gradient is suspiciously large, it is scaled down (see below).
6. The step. The learning rate is set for this step from the schedule, each optimizer moves its parameters, and the gradients are zeroed for the next round.

Gradient clipping

Sometimes a batch is unusual, the gradient comes out huge, and one careless step can throw the model far from where it was. Clipping is a seat belt: if the length (norm) of the whole gradient vector exceeds a threshold, the vector is scaled down to that length, keeping its direction:

$$g \leftarrow g \cdot \min\left(1, \frac{1}{\lVert g \rVert}\right), \qquad \lVert g \rVert = \sqrt{\textstyle\sum_i g_i^2}.$$

clip_grad_norm_ also returns the norm before clipping, and we log it. In the curves further down you can see that at the very start the norm is several times above 1, so the belt really does hold Sprout back during the first chaotic steps. Within the first hundred and fifty steps or so the norm drops below 1, and after that clipping hardly ever fires.

bf16 and torch.compile: two free speed-ups

bf16 ("brain float 16") is a 16-bit number format with the same 8 exponent bits as the usual 32-bit float but only 7 bits of mantissa. It covers the same range of magnitudes, from tiny to huge, only less precisely: two to three significant decimal digits. For training neural networks that is usually enough, and matrix multiplications in bf16 move half as many bytes and run faster. torch.autocast decides by itself which operations can go to bf16 (matrix multiplications) and which must stay in float32 (sums, softmax, the loss). The weights themselves and the optimizer's state are kept in float32.

torch.compile looks at the model's Python code once, builds a graph of operations from it and fuses small operations into larger GPU kernels, so that data is not shuttled between memory and the chip for every little addition. On our Mac this made training about twice as fast as the ordinary ("eager") mode, without a single change to the model.

How to take a step

After the backward pass we know the gradient $g$: the direction in which the loss grows fastest. The simplest step is plain gradient descent from chapter 3:

$$\theta \leftarrow \theta - \eta\, g.$$

In practice it has two weaknesses. First, the gradient of one batch is noisy: it points roughly the right way but wobbles from batch to batch. Second, a real loss landscape is shaped like a ravine: very steep across and almost flat along. A learning rate small enough not to bounce off the steep walls is far too small to make progress along the gentle floor.

Momentum treats the parameters like a heavy ball. Instead of the current gradient we step along a running sum of past gradients, $m \leftarrow \beta m + g$, $\theta \leftarrow \theta - \eta m$ with $\beta = 0.9$. The zigzags across the ravine cancel out, and the steady push along it accumulates.

Adam goes further and gives every parameter its own step size. Besides the average gradient $m$ it tracks the average square of the gradient $v$ and divides one by the other:

$$m \leftarrow \beta_1 m + (1-\beta_1)\, g, \qquad v \leftarrow \beta_2 v + (1-\beta_2)\, g^2, \qquad \theta \leftarrow \theta - \eta\, \frac{\hat m}{\sqrt{\hat v} + \epsilon}.$$

The ratio $m / \sqrt{v}$ is about $\pm 1$ when a parameter's gradient is consistent and close to zero when it keeps flipping sign. So each parameter moves by roughly $\eta$ per step, however large or small its gradient is. (The hats mark a correction for the first steps, when $m$ and $v$ are still close to their initial zeros.) AdamW adds weight decay, the penalty from chapter 3: a slight pull of every weight towards zero, applied separately from the gradient.

Let four optimizers race on three landscapes. The landscapes are two-dimensional, so we can see them, but the difficulties are the real ones: a narrow tilted ravine, a curved banana-shaped valley, and a plateau where the gradient is almost zero.

Plain gradient descent (SGD) crawls along the ravine and gets stuck on the plateau, where the gradient is tiny and its steps are tinier still. Momentum accelerates but overshoots. Adam and Muon don't care how large the gradient is, only which way it points, so they cross the plateau at full speed. That independence from scale is the main reason modern networks are trained with adaptive optimizers. But watch Muon near the goal: its step has a fixed length, so it reaches the bottom first, jumps right over it and keeps hopping around until the cooldown at the end of the schedule shrinks its steps. Adam hops too, only less. We will meet the cooldown again in the section on the schedule.

Muon: every direction gets the same step

Adam treats the parameters as a long list of independent numbers. But most of Sprout's parameters are arranged in matrices, and the update of a matrix has a structure of its own. Linear algebra has a tool for seeing it, the singular value decomposition: any matrix can be written as $G = U \Sigma V^\top$, a rotation, a stretch along a few axes by factors $\sigma_1 \ge \sigma_2 \ge \dots$, and another rotation. In the gradient of a weight matrix, typically a few directions have huge $\sigma$ and the rest have tiny ones. A plain step then moves the weights mostly along those few dominant directions, and rare but useful directions barely move.

Muon (short for MomentUm Orthogonalized by Newton–Schulz) keeps the directions and throws away the stretching: it replaces the momentum-averaged gradient with the nearest orthogonal matrix,

$$G = U \Sigma V^\top \quad\longrightarrow\quad O = U V^\top,$$

in which every singular value equals one. Every direction of the update gets the same step size.

Computing an SVD at every step would be slow. Muon uses a trick: a polynomial iteration that needs only matrix multiplications. Start from $X_0 = G / \lVert G \rVert$ (so all singular values are at most 1) and repeat

$$X \leftarrow aX + b\,(XX^\top)X + c\,(XX^\top)^2 X.$$

Here is the optimizer we trained Sprout with, the whole of muon.py minus its opening docstring and imports:

def orthogonalize(G, steps=5): """G: (..., rows, cols) -> the nearest semi-orthogonal matrices, batched.""" a, b, c = 3.4445, -4.7750, 2.0315 # tuned for fast convergence X = G.bfloat16() tall = X.size(-2) > X.size(-1) if tall: X = X.mT X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7) for _ in range(steps): A = X @ X.mT X = a * X + (b * A + c * A @ A) @ X if tall: X = X.mT return X.to(G.dtype) class Muon(torch.optim.Optimizer): def __init__(self, params, lr=0.02, momentum=0.95, weight_decay=0.0, split=None): # split: {param: n} treats a fused weight (like q, k, v) as n separate matrices self.split = split or {} super().__init__(params, dict(lr=lr, momentum=momentum, weight_decay=weight_decay)) @torch.no_grad() def step(self): for group in self.param_groups: # matrices of the same shape are orthogonalised together, in one batch by_shape = defaultdict(list) for p in group['params']: if p.grad is None: continue state = self.state[p] if 'momentum' not in state: state['momentum'] = torch.zeros_like(p) buf = state['momentum'] buf.lerp_(p.grad, 1 - group['momentum']) g = p.grad.lerp(buf, group['momentum']) # Nesterov look-ahead n = self.split.get(p, 1) by_shape[(p.size(0) // n, p.size(1))].append((p, g.view(n, -1, p.size(1)))) for (rows, cols), items in by_shape.items(): updates = orthogonalize(torch.cat([g for _, g in items])) scale = max(1.0, rows / cols) ** 0.5 i = 0 for p, g in items: u = updates[i:i + g.size(0)].reshape_as(p) i += g.size(0) if group['weight_decay']: p.mul_(1 - group['lr'] * group['weight_decay']) p.add_(u, alpha=-group['lr'] * scale)

A few details that matter in practice. The momentum uses a Nesterov look-ahead: the update mixes the current gradient with the momentum buffer. Matrices of the same shape are orthogonalised together in one batch, which is why Muon costs so little. The fused matrix that computes queries, keys and values at once is first split into its three parts, since in meaning they are three separate matrices. Only the two-dimensional weights inside the blocks go to Muon; the embedding table and the RMSNorm gains stay with AdamW.

You can check the iteration yourself in numpy: take a "gradient" in which one direction outweighs another a hundredfold and orthogonalise it.

import numpy as np def orthogonalize(G, steps=5): a, b, c = 3.4445, -4.7750, 2.0315 # the same coefficients as muon.py X = G / (np.linalg.norm(G) + 1e-7) for _ in range(steps): A = X @ X.T X = a * X + (b * A + c * A @ A) @ X return X rng = np.random.default_rng(0) # a "gradient" of a 4x6 weight matrix where a few directions dominate G = rng.normal(size=(4, 6)) @ np.diag([10, 3, 1, 0.3, 0.1, 0.03]) U, S, Vt = np.linalg.svd(G, full_matrices=False) print('singular values of G: ', S.round(3)) print('after orthogonalize(G): ', np.linalg.svd(orthogonalize(G), compute_uv=False).round(3)) print('exact U @ Vt: ', np.linalg.svd(U @ Vt, compute_uv=False).round(3))

Our A/B

An optimizer is not a matter of taste; you measure it. Before the main run we trained the same Sprout three times on the same 10 million tokens: with AdamW and with Muon at two learning rates.

The result is unambiguous: 2.393 for Muon against 2.709 for AdamW on validation. A difference of 0.32 nats means that Muon's model is choosing among $e^{0.32} \approx 1.37$ times fewer options per token. Muon is a little slower per step (29.7 thousand tokens per second against 31.8 thousand), but on quality per token it wins by a mile. Doubling its learning rate to 0.04 changed almost nothing, a sign that it is not fragile. That settled it: the main run uses Muon with lr 0.02.

What does Muon do with the update of a weight matrix before applying it?

Divides each element by the root of its average squared gradient Keeps the directions of the update but makes all of its singular values roughly equal Throws away the small elements of the gradient and keeps the large ones

That is orthogonalisation: $U\Sigma V^\top \to UV^\top$. Dividing element by element is what Adam does; Muon looks at the matrix as a whole and evens out its directions, so rare directions are not drowned out by dominant ones.

The learning-rate schedule

The learning rate $\eta$ does not stay constant during training. Our schedule has three parts:

• Warm-up, 200 steps: the rate grows linearly from almost zero to its peak. At the start the weights are random, the gradients are large and chaotic, and the optimizers' averages ($m$, $v$ and the Muon buffer) are still empty. A full-size step at that moment can throw the model into a region it will take thousands of steps to leave.
• Hold: the peak rate for most of the run. Large steps move fast, but the noise of the batches keeps the model rattling around the bottom of the valley instead of settling into it.
• Cooldown, the last 30%: the rate falls linearly to zero. The steps get smaller, the noise averages out, and the model settles into the bottom. The loss usually drops noticeably during this phase.

Try it on a toy problem: a noisy valley in 24 dimensions with directions of very different steepness. Far from the bottom its walls are eight times steeper, just as a real network's landscape is sharpest at the start of training.

A few experiments worth doing. Remove the warm-up: the very first steps throw the ball off the cliff. Remove the cooldown ("Constant"): the loss stays at the level of the noise. Raise the peak: at first progress is faster, but past some point the noise wins, and a little further on even a warm-up cannot save you. Compare the linear cooldown with a cosine: the final results are close, which matches experience with real models.

Here are the two functions from train.py that set the schedule, unchanged, and what they give for Sprout's run:

def lr_factor(step, total, warmup, cooldown): """Warm up, hold, then decay linearly to zero over the last `cooldown` share.""" if step < warmup: return (step + 1) / warmup decay_start = total * (1 - cooldown) if step < decay_start: return 1.0 return max(0.0, (total - step) / (total - decay_start)) def context_at(step, total, context, warm): """Sequence-length warm-up: short windows first, full length later.""" if not warm: return context if step < 0.1 * total: return max(64, context // 4) if step < 0.3 * total: return context // 2 return context total = int(330e6 // 32768) # 10,070 steps for step in [0, 99, 199, 1006, 1007, 3021, 7048, 8500, 10069]: T = context_at(step, total, 512, True) print(f'step {step:5d} lr x {lr_factor(step, total, 200, 0.3):.3f} window {T} windows per batch {32768 // T}')

The second function is the sequence-length warm-up. For the first 10% of steps Sprout reads windows of 128 tokens, then windows of 256 until the 30% mark, and only then the full 512. The number of tokens per batch stays the same; there are just more, shorter windows. Short windows are cheaper, since the cost of attention grows as the square of the length, and early on the model is busy with local patterns anyway: words, punctuation, grammar within a sentence. Long-range connections are what it learns later, when it is ready for them.

The real run

Now the main run: 10,070 steps, 330 million tokens, Muon and AdamW, the schedule above, all on one Mac with an M4 Pro. It took 3 hours 41 minutes (13,289 seconds), 24.8 thousand tokens per second on average, evaluations included, and ended at a validation loss of 1.554 nats per token. Every 20 steps the script logged the loss of the batch, the learning rate, the window length and the gradient norm; every 250 steps it measured the loss on the validation text (the same 20 batches of 512-token windows every time).

What to look for:

• The first hundred steps take the loss from about 9 to below 4. That is the model discovering the frequencies of tokens and the simplest pairs. Switch to the log scale: on it the whole run is easier to see, and the long middle part becomes a gentle, almost straight slope.
• The training and validation curves move together. The model almost never sees the same text twice (remember the mosaic from the last chapter: four fifths of the corpus stays unread), so it has no chance to memorise its training data. In this regime there is no overfitting, and the validation loss is an honest measure of progress.
• The window switches leave a trace. Validation is always measured on 512-token windows, while for the first 10% of training Sprout has never read anything longer than 128 tokens. The exam is harder than the lessons, and right after the switch to longer windows the validation loss drops noticeably faster. The training loss also jumps down at each switch: with a longer context the next token is easier to guess.
• The gradient norm (lower panel) starts well above the clipping threshold and falls below it within the first hundred and fifty or so steps.

On the log scale the middle of the curve is a long, gentle slope: every tenfold increase in steps takes off a similar chunk of the loss, a little smaller each time. Curves like this follow a power law, $\mathcal{L}(D) \approx E + A \cdot D^{-\alpha}$, where $E$ is the part of the loss that no amount of data can remove. Each tenfold increase in data removes the same share of what remains above $E$; on log–log axes, $\mathcal{L} - E$ is a straight line. This is the scaling law that Kaplan et al. (2020) found across models of many sizes and that the Chinchilla paper refined in 2022. It is the reason people can predict the loss of a huge model from a series of small ones before spending millions on training it. We return to it in the last chapter.

Why does Sprout's validation loss follow its training loss so closely, with no sign of overfitting?

Because the validation text was taken from the training text Because in 330 million tokens Sprout reads less than a quarter of its corpus and almost never sees the same text twice Because gradient clipping prevents overfitting

Overfitting means memorising specific examples, and for that you have to see them repeatedly. Sprout makes less than one pass over its data, so almost every batch is new text for it, just like the validation set. Models trained for many epochs on a small dataset are a different story.

Checkpoints side by side

During the run we saved Sprout's weights at several moments. Give them the same prompt and the same random seed and compare. Each checkpoint is a full 17.5 MB model, so the first time round they download one after another.

The ladder at the bottom puts Sprout in the context of the whole course. The letter bigram from chapter 1, the letter networks, the token network, the one- and four-block transformers: each one lowered the loss a notch. Sprout at step 200, after only 6.5 million tokens, is already ahead of the token network. By step 1,000 it has all but caught up with the four-block transformer of chapter 8, and the rest of the run takes it much further: compare the validation loss at step 1,000 and at the end on the curves above. The ideas were the same all along; what it took was size, data and a well-built training loop.

Sprout right now

This is the result of pre-training: a base model. It has read a third of a billion tokens, and it continues any text in the style of its corpus. It writes children's stories with named characters, keeps dialogue in the Name: line format and knows how a story begins and ends. But it does not answer questions: ask it something and it will simply carry on writing, as if your question were a line from a story or a play, with more lines, a narrator and invented speakers, and no answer addressed to you. And it picks every word with a roll of dice that we set up ourselves. In the next chapter we take the dice apart: temperature, top-k, top-p, and why greedy decoding gets stuck in loops.