Training

Sprout has an architecture and a prepared library, but its weights are still random. Now training begins. At each step it receives passages of text, predicts the next token at every position, and compares those predictions with the actual continuations. The backward pass finds how the loss depends on the parameters; the optimizer then updates them. We will repeat this cycle 10,070 times and watch the writing change in samples saved along the way.

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: its random preferences do not yet reflect patterns in the text. 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 too. "Choices per token" is a visual shorthand for perplexity, $e^{\text{loss}}$, from chapter 2. It does not count the candidates left in the vocabulary: it expresses average prediction loss on the scale of guessing among equally likely alternatives. Initially it is slightly above 8,192, since random preferences can be worse than a uniform prediction. By the end it is down to a few units. That is a substantial gain in next-token prediction, though this metric alone cannot tell us whether a whole story is truthful or coherent.

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: speeding up computation

bf16 (bfloat16) is a 16-bit format with eight exponent bits, like float32, and seven explicitly stored fraction bits. Its range of normal numbers is close to float32's, but its precision is lower: roughly two to three significant decimal digits. Mixed-precision training runs some operations, particularly matrix multiplications, in bf16 and keeps others in float32. The rules used by torch.autocast depend on the operation and device. This can reduce data movement and use fast hardware while retaining more accurate calculations where needed. Our run keeps the weights and optimizer state in float32.

torch.compile captures model computations in a graph and tries to execute them more efficiently, for instance by fusing several operations into one GPU kernel. Fewer intermediate writes to memory mean less work. Compilation itself takes time, and new input shapes can trigger recompilation. Our Mac measurements showed roughly twice the speed of ordinary eager execution; the gain on another device or model needs to be measured again.

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}.$$

If a gradient keeps roughly the same size and sign for long enough, $\hat m / \sqrt{\hat v}$ approaches $+1$ or $-1$, ignoring $\epsilon$. If the sign often changes, the average gradient can partly cancel out, reducing the update. This is a useful intuition, not an exact description of every step: the two averages use different memories and depend on the history. The hats correct for their initial zeros. AdamW applies weight decay separately, slightly shrinking the weights without mixing that shrinkage into the adaptively scaled 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.

In this demonstration, SGD moves slowly across flat regions: a small gradient means a small step. Momentum helps it accelerate but can carry it past the target. Adam scales coordinates using their gradient history; Muon here approximately normalises a one-row matrix, so a small raw gradient need not mean a small update. Their paths still depend on the settings. Watch how large steps near the minimum turn progress into repeated overshoots. Reducing the learning rate can help the run finish more carefully; we will return to that below.

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 (MomentUm Orthogonalized by Newton–Schulz) treats the momentum-smoothed update as a matrix. Its idealised goal is to remove scale differences between singular directions by replacing nonzero singular values with ones. For a full-rank square matrix this gives the nearest orthogonal matrix; for a rectangular matrix, its rows or columns are orthonormal:

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

For a full-rank matrix, this gives unit singular values. The actual iteration below only approximates that equalisation; it does not create directions whose singular values are exactly zero.

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.

In these short runs, Muon reached a validation loss of 2.393, compared with AdamW's 2.709. The gap of about 0.32 nats corresponds to a perplexity ratio of $e^{0.32} \approx 1.37$. Muon was slightly slower: 29.7 thousand tokens per second against 31.8 thousand. Its better result for the same token budget led us to choose it for the main run at a learning rate of 0.02. The result at 0.04 was similar. These three experiments were enough for our practical choice, but do not establish that one optimizer always wins.

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 singular directions and approximately equalises nonzero singular values 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 building a schedule for a toy problem: a noisy valley in 24 dimensions. Its directions are deliberately given different slopes, and the walls become eight times steeper far from the bottom. This construction helps illustrate why a cautious start and small finishing steps can help. It models a few optimisation difficulties, not the full loss landscape of a language model.

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 sets the sequence-length warm-up: 128-token windows for the first 10% of steps, 256 until the 30% mark, then 512. Each batch still contains 32,768 tokens, with more windows when they are shorter. Attention pairs grow quadratically with the length of a single window, but at a fixed token count per batch the total grows linearly with window length. Short windows therefore save part of the computation. Sprout first trains on local connections; extending the window then makes more distant connections available to learn from.

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. Most of the corpus remains unread, with relatively few repeated presentations of particular passages. This supplies varied examples. The curves show no clear divergence, although that does not prove an absence of memorisation.
• 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.

What helps Sprout keep its training and held-out losses close during this run?

Because the validation text was taken from the training text A large corpus provides many new examples and relatively few repeated presentations Because gradient clipping prevents overfitting

Overfitting occurs when improvements on the training set stop carrying over to new examples. A large, varied corpus with infrequent repetition helps reduce this risk, but does not rule it out or prevent memorisation. Our curves show no clear divergence during this run. It is that observation on held-out text, rather than the token count alone, that lets us assess progress.

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 has not been trained to answer as an assistant: ask it something and it may 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.