A bigger model

Sprout-1 was 8 blocks of 384 numbers with 6 attention heads, a vocabulary of 8,192 tokens and a window of 512. Sprout-2 is 30 blocks of 576, with 9 query heads sharing 3 key/value heads, a vocabulary of 32,768 and a window of 2,048: 125.1 million parameters instead of 17.3 million. Sprout-1 read 330 million tokens in 10,070 steps, three and a half hours on a Mac. Sprout-2 read 9.0 billion, in 17,166 steps of 524,288 tokens each.

Nothing from chapters 7, 8 and 10 goes out of date: the same attention, the same block of attention plus SwiGLU, the same Muon and AdamW, the same loss. What changes when a model grows is quieter: its shape, what fills its memory while it writes, a few safety belts that a small model could do without, and how to choose all of that before the long run. But first, the result.

Watch it grow

At eleven moments of its training Sprout-2 was given the same five prompts: the start of a story, the start of an explanation, a Python function, a shell command and a question. Pick a prompt and climb the ladder, or press "Grow".

A few things to look for.

• Step 200, 105 million tokens, a third of everything Sprout-1 ever read. The story already has the rhythm of a children's tale. The function has its indents in the right places and computes nothing: weight = 0, height = 1. The command turns into a pipeline of xargs and awk.
• By step 2,000 the story holds together: a rabbit, a log, a wish to see what is inside. By step 4,000 the tar command is followed by what a terminal would print next: ls -l and a listing with permissions and owners.
• The function takes longest. At step 8,000 it has a proper docstring, "Return the n-th fibonacci number.", and then returns 2, 3, 4 and 5. At step 16,000 it writes the textbook recursion, fibonacci(n-1) + fibonacci(n-2), with a base case that would never let it finish. At the very last step it invents "the FIBONAREA". A single sample is one roll of the dice; the curve beside it is the measurement.
• The question gets no answer at any step. From step 6,000 on, Sprout-2 replies the way a forum thread goes on: "Here is a list of all files that are hidden from view. How do I get the list of hidden files…?" At step 16,000 it even offers a find command, followed, true to the forums, by the error message someone got from it. Remember this; we'll come back to it at the end.

Now let's open it up.

Where the parameters go

Let's take Sprout-2 apart the way chapter 8 took Sprout-1 apart. Side by side:

	Sprout-1	Sprout-2
blocks	8	30
stream width $d$	384	576
query heads of 64	6	9
key/value heads	6	3
SwiGLU width $d_{ff}$	1,024	1,536
vocabulary	8,192	32,768
window	512	2,048
RoPE base	10,000	100,000
logits	as they are	soft-capped at 30
parameters	17,309,056	125,081,664

Count one block. Attention has four matrices, but only two of them are square: the queries $W_Q$ and the output projection $W_O$ are $576 \times 576$, while keys and values need only 3 heads of 64, so $W_K$ and $W_V$ are $576 \times 192$. The SwiGLU network has three matrices of $576 \times 1{,}536$. Then come two norms and the QK-norm gains:

$$2 \cdot 576^2 + 2 \cdot 576 \cdot 192 + 3 \cdot 576 \cdot 1{,}536 + 2 \cdot 576 + 2 \cdot 64 = 3{,}540{,}224.$$

Thirty such blocks, the embedding table and the final norm:

$$32{,}768 \cdot 576 + 30 \cdot 3{,}540{,}224 + 576 = 125{,}081{,}664.$$

Sprout-2 is 7.2 times bigger than Sprout-1, and its blocks 7.5 times: the growth went into the parts that think. The SwiGLU networks alone hold 79.6 million numbers, 64% of the model. The embedding table grew four times in rows and one and a half times in width, to 18.9 million numbers, yet its share fell from 18% to 15%. And it is still one table used at both ends (chapter 8): untied, the model would carry another 18.9 million.

Two lines of the config cost no parameters. The window is 2,048 tokens, four times Sprout-1's, and RoPE's base is 100,000 instead of 10,000. A larger base slows RoPE's slower rotations (chapter 8), the slowest almost ten times, so that positions far apart in a longer window still look different: the usual adjustment when a window grows.

Deep and thin

Why 30 blocks of 576 and not, say, 16 of 768, which holds almost exactly as many parameters? Switch the widths in the lower half of the widget.

A wide model spends more of its budget on the vocabulary: at a width of 1,152 the embedding table eats 31% of the model, and only 6 blocks are left. A deep model gives every token more steps of processing, one after another, and chapter 8 showed why steps matter: an induction head needs an earlier layer to prepare what it reads. Depth has a price too: thirty layers must run one after the other, which suits a GPU less than a few wide ones, and the KV cache of the next section grows with every layer.

Here is the count as a function you can run, with a third model: Sprout-2 without its grouped-query attention.

def count(vocab, d, layers, heads, kv_heads, d_ff): hd = d // heads # 64 numbers per head attn = 2 * d * heads * hd + 2 * d * kv_heads * hd # q and proj; the narrower k and v ffn = 3 * d * d_ff # SwiGLU: w1, w3 and w2 norms = 2 * d + 2 * hd # two RMSNorms, the QK-norm gains emb = vocab * d # one table, used at both ends total = emb + layers * (attn + ffn + norms) + d # + the final norm return total, emb, layers * attn, layers * ffn for name, cfg in [('Sprout-1', (8192, 384, 8, 6, 6, 1024)), ('Sprout-2', (32768, 576, 30, 9, 3, 1536)), ('Sprout-2, 9 k/v heads', (32768, 576, 30, 9, 9, 1536))]: total, emb, attn, ffn = count(*cfg) print(f'{name:21} {total:>11,} embeddings {emb / total:.0%}' f' attention {attn / total:.0%} SwiGLU {ffn / total:.0%}')

With as many key/value heads as query heads, Sprout-2 would weigh 138.4 million: 13.3 million more, all of them in $W_K$ and $W_V$. That is the smaller of the two savings. The bigger one shows only when the model writes.

Fewer keys than queries

Chapter 11 introduced the KV cache: while the model writes, it keeps the key and the value of every past token in every layer, so each new token costs a single pass. The price is memory, and it is easy to count. For every token the model stores

$$2 \times \text{layers} \times \text{key/value heads} \times \text{head size}$$

numbers: a key and a value in every layer, for every head that has keys and values of its own. For Sprout-1 that was $2 \times 8 \times 6 \times 64 = 6{,}144$ numbers. Give Sprout-2 as many key/value heads as query heads and it would be $2 \times 30 \times 9 \times 64 = 34{,}560$ numbers per token: 69 KB in 16-bit numbers, 141.6 MB for a full window of 2,048 tokens, almost twice the 78 MB of its 4-bit weights.

Grouped-query attention (GQA) shares keys and values. Sprout-2's 9 query heads form 3 groups of three, and each group reads one set of keys and values. Every head still asks its own question with its own slice of $W_Q$; only the labels it compares against and the contents it mixes (chapter 7) are shared by the three heads of a group. The cache shrinks three times, to 11,520 numbers per token.

In the browser Sprout-2's engine keeps keys and values as 32-bit floats, so a full 2,048-token window takes 94.4 MB, more than the 4-bit weights themselves. With 9 key/value heads it would be 283 MB. The engine sets aside at most 256 MB for caches: enough for two full conversations with 3 heads, not enough for even one with 9. And the weights are loaded once, while every conversation needs its own cache, which is why servers that talk to thousands of people at once care about the cache even more than a phone does.

In model.py the whole idea is a few lines of the Attention class:

class Attention(nn.Module): def __init__(self, cfg): super().__init__() self.n_head, self.n_kv, self.hd = cfg.n_head, cfg.n_kv_head, cfg.head_dim self.q = nn.Linear(cfg.d_model, cfg.n_head * self.hd, bias=False) self.k = nn.Linear(cfg.d_model, cfg.n_kv_head * self.hd, bias=False) self.v = nn.Linear(cfg.d_model, cfg.n_kv_head * self.hd, bias=False) self.proj = nn.Linear(cfg.n_head * self.hd, cfg.d_model, bias=False) self.q_norm = RMSNorm(self.hd) self.k_norm = RMSNorm(self.hd) def forward(self, x, cos, sin, mask=None, cache=None): B, T, _ = x.shape # (B, T, C) -> (B, heads, T, head_dim) q = self.q(x).view(B, T, self.n_head, self.hd).transpose(1, 2) k = self.k(x).view(B, T, self.n_kv, self.hd).transpose(1, 2) v = self.v(x).view(B, T, self.n_kv, self.hd).transpose(1, 2) q, k = self.q_norm(q), self.k_norm(k) q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin) if cache is not None: # generation: remember keys and values, attend to all of them k, v = cache.add(k, v) if mask is not None and not torch.is_tensor(mask): # a FlexAttention block mask (CUDA) y = flex_attention(q, k, v, block_mask=mask, enable_gqa=True) else: rep = self.n_head // self.n_kv # every key/value head serves `rep` query heads k, v = k.repeat_interleave(rep, dim=1), v.repeat_interleave(rep, dim=1) # softmax(q·k / sqrt(d)) · v, every token sees only itself and the past if mask is None: y = F.scaled_dot_product_attention(q, k, v, is_causal=True) else: y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) y = y.transpose(1, 2).contiguous().view(B, T, -1) return self.proj(y)

Compared with Sprout-1, the fused qkv matrix is gone: k and v now have their own, narrower layers, 192 outputs instead of 576. Look at the order of the lines: the cache takes the keys and values before repeat_interleave copies each of them for its three query heads, so only three heads are ever stored. The copying makes the rest ordinary attention from chapter 7, and on a GPU FlexAttention skips even the copies (enable_gqa=True).

Sprout-2 has 9 query heads and 3 key/value heads. What would change with 9 key/value heads?

It could look back at three times as many tokens Its KV cache would be three times larger, and the model 13.3 million parameters heavier Nothing: the cache keeps queries, and there would still be nine of them

The cache keeps keys and values, never queries: a query is used once, by the token that asks it. The window stays 2,048 tokens either way. What grows is the memory per token, from 11,520 to 34,560 numbers, and the $W_K$ and $W_V$ matrices in every block.

What a bigger model needs

Some habits a small model gets away with stop being harmless when training runs for billions of tokens, and some protections cost almost nothing. Sprout-2 adds two and keeps one from Sprout-1.

One row, several documents

Training reads rows of 1,024 tokens (2,048 at the end). Documents don't come in that size: an average TinyGSM problem is 216 tokens long, a TinyStories tale 196, a FineWeb-Edu page 1,070. So prepare.py writes every source as one long stream of tokens, each document opened by <|endoftext|>, and train.py cuts rows out of the stream wherever they happen to fall. A row usually starts in the middle of one document and holds the beginnings of others.

With the causal mask alone (chapter 7), every token may look at everything before it in the row, including documents that have nothing to do with it. Here is one packed row: the end of a story, a shell command and the start of a web page.

Switch to "Causal mask only" and tap "·ls": 6 of the 9 tokens it may look at are the end of a fairy tale. The model can learn to ignore such noise, but it costs effort, and short documents make it worse: in a 1,024-token row of TinyGSM problems, about 80% of the pairs that causal masking allows cross into another problem.

The fix is to number the documents and let attention through only within one. train.py numbers them in one line, (x == eot).cumsum(1): a running count of <|endoftext|> markers, so each marker opens a new document; model.py turns the numbers into a mask (the code is at the end of this section). On the small test models the run without the mask ended at 3.367 nats and with it at 3.342: 0.024 better, about what separates a good learning rate from one twice too large. And it's free: with torch.compile on the Mac the mask didn't change the speed, and on a GPU FlexAttention skips the parts of the matrix that are masked entirely.

A ceiling for the logits

At the top of the model 32,768 logits go into softmax (chapter 2). Nothing limits their size, and training has a reason to inflate them: a bigger logit on an easy token means a surer model and a lower loss. A runaway logit makes the model overconfident, and then one confident mistake costs a huge loss. Sprout-2 puts a soft ceiling over every logit:

$$\ell' = 30 \tanh\!\left(\frac{\ell}{30}\right).$$

Near zero, tanh is almost a straight line: 5 becomes 4.95, 10 becomes 9.65, so ordinary logits pass nearly untouched. Past 30 the curve bends: 60 becomes 28.9, 100 becomes 29.9, and no logit ever leaves the band from −30 to 30. Two things follow. The slope, $1 - \tanh^2(\ell/30)$, is 0.42 at 30 and 0.07 at 60: the further a logit runs, the less training can push it, so the reason to inflate logits disappears. And the cost of any single token is bounded: even the most confident mistake can't cost more than $2 \cdot 30 + \ln 32{,}768 \approx 70.4$ nats.

On the small test models the cap helped a little: 3.351 without it, 3.342 with it. At the first check the uncapped model was even slightly ahead; the cap paid off by the end. The engine in your browser applies the same function to every logit, so Sprout-2 speaks there exactly as it was trained.

QK-norm stays

The other place where numbers can run away is attention itself: $q \cdot k$ grows with the length of the vectors. Sprout-1 already normalised every query and key before comparing them (QK-norm, chapter 7), and Sprout-2 keeps it. Between them, QK-norm and the soft-cap guard both of the model's softmaxes: the one inside attention and the one at the output. Here are the mask and the cap in model.py:

def doc_mask(docs): """Dense (batch, 1, T, T) mask: a token sees the past of its own document only.""" T = docs.size(1) causal = torch.ones(T, T, dtype=torch.bool, device=docs.device).tril() return (causal & (docs[:, :, None] == docs[:, None, :]))[:, None] # ... def logits_of(self, x): logits = self.head(self.norm(x)).float() if self.cfg.softcap: logits = self.cfg.softcap * torch.tanh(logits / self.cfg.softcap) return logits

The mask is exactly the widget's matrix: the causal triangle (tril) AND "same document number". The cap is one line, applied after the output layer and before the loss.

Rehearsals on small copies

How do you know that the mask gains 0.024, or that 0.02 is the right learning rate for Muon, before the long run? Not by asking the big model. You build a small copy of it, a proxy, and change one thing at a time.

Sprout-2's proxies have 8 blocks of 256, with 4 query heads and 2 key/value heads: 14.3 million parameters, almost 60% of them in the embedding table, which stays 32,768 tokens tall. Each proxy read 100 million tokens from a 1.4-billion-token miniature of the corpus, about an hour on the Mac; everything else follows Sprout-2's recipe, scaled down. Here are all the curves.

Muon's learning rate of 0.02 ends at 3.342, ahead of 0.01 (3.353) and 0.04 (3.368). For the embedding table, trained by AdamW, 0.006 beats both 0.003 (3.361) and 0.012 (3.347). The document mask gains 0.024 and the soft-cap 0.008. Now switch to "Raw loss": the curves lie almost on top of each other. Every decision here comes down to hundredths and thousandths of a nat on a curve that falls from 4.8 to 3.3. Early gaps are bigger (at the first check Muon at 0.04 is 0.18 behind) and mostly shrink; what counts is the end. The sharp drop after 80 million tokens is the proxies' own cooldown.

Two cautions. Each variant was trained once, so we don't know how far a different random seed would move its result: the thousandths deserve more doubt than the hundredths. And a proxy is nine times smaller than Sprout-2 and reads ninety times less, so values that are best for it need not be best for the big model. Sprout-2 simply took them.

The schedule

Chapter 10's schedule had three parts: warm-up, hold, cooldown. Sprout-2 keeps the shape and changes the proportions:

• Warm-up, 250 steps, 131 million tokens: the learning rate grows from almost zero to its peak.
• Hold at the peak until step 13,732.
• Cooldown, the last 20%: 3,434 steps and 1.8 billion tokens during which the rate falls linearly to zero. The mix changes too (chapter 14): more code, questions and answers and maths, less web, and 3% conversations in the chat format.

The rows grow along the way: 256 tokens for the first 2% of the steps, 512 until 6%, 1,024 for most of the run and 2,048 only for the last 10%. The batch is always 524,288 tokens, so it holds 2,048 rows of 256 at the start and 256 rows of 2,048 at the end. Long rows come last because the cost of attention grows with the square of the length (chapter 7), and because a long past is of most use to a model that already writes well. Here are the two functions from train.py that set all this:

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 = int(total * (1 - cooldown)) if step < decay_start: return 1.0 return max(0.0, (total - step) / (total - decay_start)) def context_at(step, total, args): """Short rows first; the full --long-context only in the last `long_frac` of the run.""" frac = step / max(1, total) if args.seq_warmup and frac < 0.02: return args.context // 4 if args.seq_warmup and frac < 0.06: return args.context // 2 if frac >= 1 - args.long_frac: return args.long_context return args.context

For 17,166 steps they give rows of 256 tokens up to step 343, of 512 up to step 1,029, of 1,024 up to step 15,449 and of 2,048 from step 15,450; the cooldown starts at step 13,732. Sprout-1's version of context_at grew the window at the start only; Sprout-2's also saves its longest rows for the very end.

A checkpoint before the cooldown

At step 13,732, right before the rate starts to fall, train.py saves a full checkpoint, ckpt_pre_cooldown.pt. This is the practical gift of the warm-up, hold and cooldown shape (chapter 10): at that moment the model is in the middle of its training, not at its end. If more data turns up later, training can continue from that checkpoint at full speed, as if the cooldown had never happened, and cool down again at the new end. The finished model is not a dead end.

What a full checkpoint holds

A final checkpoint needs only the weights. The one before the cooldown also holds the state of both optimizers (Muon's momentum, AdamW's running averages), the loader's place in every source (which one-million-token chunk, in which shuffled order, how far into it) and the step and token counts. With --resume all of it comes back, and the continued run reads the very tokens it would have read next; only the learning rates are set anew, so a continuation may choose its own.

The run

How much work is this? Training costs about six operations per parameter per token, $6ND$ (chapter 9): two for the forward pass and four for the backward. For Sprout-1 that is $6 \times 17.3 \cdot 10^6 \times 330 \cdot 10^6 \approx 3.4 \cdot 10^{16}$; for Sprout-2, $6 \times 125.1 \cdot 10^6 \times 9.0 \cdot 10^9 \approx 6.75 \cdot 10^{18}$, about 200 times more. That is why Sprout-2 couldn't grow on a laptop the way Sprout-1 did: it trained on one datacenter GPU of the H100 class, for about 9.5 hours.

Every 500 steps train.py measured the loss on held-out text from each of the 31 sources and the conversations. Here is the whole run.

What the curves say:

• Every group keeps improving to the end. Web text falls from 3.75 at step 500 to 2.94, code from 2.26 to 1.47, maths from 1.98 to 1.35. The curves flatten but keep falling: 9 billion tokens, 72 per parameter, did not exhaust what a model of this size can learn.
• Conversations wait for the cooldown. Their loss falls with ordinary English at first and then stays near 2.6. At step 13,732 the cooldown mix brings chat data in; by step 14,000 the loss is 1.59, and it ends at 1.43. Those 268 steps held about 4.2 million tokens of conversations: the words were not new to the model, the format was.
• The training loss jumps down at the start of the cooldown, from about 2.51 to about 2.23 within a hundred steps, while the learning rate is still at 98% of its peak. That is not the model settling: the mix has changed, and code, maths and chats are easier to predict than web pages. The black line, one fixed exam, barely moves: 2.581 at step 13,500, 2.578 at step 14,000, and 2.505 at the end of the cooldown.
• Switch on "val_mix as logged". That is the number train.py prints, and it drops from 2.58 at step 13,500 to 2.30 at step 14,000. val_mix weighs the sources by the current mix, so at step 13,732 the exam itself changed: easier subjects got more weight, and the conversations joined in.
• Longer rows help a little. When rows grow to 2,048 tokens at step 15,450, the training loss steps down again, from 2.21 to 2.12 averaged over a thousand steps on each side: with a longer past there is more to predict from, as chapter 10 saw with Sprout-1.

Now switch to "By source". At the end TinyGSM's problems cost 0.64 nats per token and DCLM's web pages 3.37. Sprout-2 did not learn maths five times better than English: a TinyGSM problem is a short word problem with a Python solution in one fixed style, like two million others, while a web page can be about anything. Loss measures how predictable a kind of text is as much as how well it was learned, so compare losses only on the same text. Not even TinyStories' 1.34 against Sprout-1's 1.554 (chapter 10): the two models cut text into different tokens, and a loss per token depends on what a token is.

Between steps 13,500 and 14,000 the val_mix that train.py logs fell from 2.58 to 2.30. What was the main reason?

The falling learning rate let the model settle into the bottom of the valley From step 13,732 val_mix weighs the sources by the cooldown mix and includes the conversations, so the exam itself changed The rows grew to 2,048 tokens, giving the model more context

On one fixed exam the loss went only from 2.581 to 2.578 over those steps. The rest of the drop is the change of weights: code and maths, which have low losses, count for more in the cooldown mix, web pages for less, and the conversations, which the model learns fast, are counted at all. The rows grow to 2,048 only at step 15,450.

The base model

After 17,166 steps we have a base model, one that has only ever learned to continue text. How much does it know? A base model is asked to choose, not to write: every benchmark question comes with a few possible endings, and eval.py picks the ending the model finds likeliest after the question. acc compares the plain sums of log-probabilities; acc_norm divides each by the ending's length in characters, so long endings are not punished for being long. Every item of every benchmark was used.

Benchmark	Chance	acc	acc_norm
SciQ a science question after a supporting paragraph	0.25	0.799	0.746
ARC-Easy school science questions	0.25	0.506	0.459
ARC-Challenge the questions simple methods get wrong	0.25	0.219	0.250
HellaSwag the likeliest continuation of an everyday scene	0.25	0.305	0.341
PIQA which of two ways to do something physically works	0.5	0.645	0.638
WinoGrande who or what a pronoun refers to	0.5	0.520	0.520
OpenBookQA applying a science fact to a new situation	0.25	0.180	0.304
BoolQ yes or no about a paragraph	0.5	0.579	0.619

Read it against the chance column. SciQ, ARC-Easy, PIQA and HellaSwag are clearly above it: with a paragraph to read from, Sprout-2 picks the right science answer four times out of five, and it has some everyday sense of what follows what. ARC-Challenge and WinoGrande sit at chance: questions built to defeat word association, and pronouns that need real understanding, are beyond it. BoolQ looks above chance but isn't: 2,033 of its 3,270 answers are "yes", so always saying yes would score 0.622.

Writing is harder than choosing. On GSM8K's word problems, written as the start of a TinyGSM-style Python function and scored by running the program it writes, Sprout-2 gets 21.5% right: the style of TinyGSM from its maths mix. On MBPP and HumanEval, which ask for a whole function that passes tests, it scores 1.0% and 2.4%. With greedy decoding it falls into loops, writing the same import twenty times or copying a function it has just written. The test harness itself was checked on the reference solutions: 40 out of 40 passed.

And the question from the top of the chapter? To "How do I list all files in a folder, including hidden ones?" the finished model replies: "Here is a list of all files that are hidden from view. How do I get the list of hidden files (with 0 size, 0 characters) when the folder is listed in the list in the admin console?" It doesn't answer; it continues a forum thread, because on the web a question like that is usually followed by more questions.

Sprout right now

Here is Sprout-2: the finished base model from the top rung of the ladder. It answers from our server, which runs the same engine your browser runs for the course's other models, on the same file: 4 bits per weight, 78 MB instead of the 250 MB the weights take in 16 bits, for a price of 0.026 nats of loss. Nothing is downloaded, and nothing is stored. Try the five prompts from the start of the chapter: with fresh dice they come out differently every time. It writes stories, explanations, Python and shell sessions, and it still turns every question into a forum thread: it continues text, it doesn't answer. The next chapter teaches it to talk, to think before it answers and to run Python when arithmetic gets hard.