Chapter 16 of 17 35 min
Thinking and a Python tool
The base model continues text. A chat format with more roles, a quarter of a million conversations and four rounds of fixes turn it into a chat that thinks before it answers and hands arithmetic to a Python program that really runs.
In this chapter
- read a chat format with a system message, thoughts and tool calls token by token, and say which tokens the loss learns
- explain why writing out steps helps a small model, where Sprout-2's thoughts come from, and why a tool beats thoughts at arithmetic
- follow fine-tuning as a series of fixes: a symptom, the data that caused or cured it, and a measurement
Chapter 15 ended with a model that knows a great deal about what text looks like and answers a question only when an answer happens to be the likeliest continuation. This chapter turns it into the chat at the bottom of the page. That takes a chat format with more roles than Sprout-1 had, thoughts before answers, a Python tool that really runs, and four rounds of fine-tuning.
Nobody planned four rounds: each later one fixes something the one before it broke or never taught. That is what fine-tuning usually looks like from the inside, and this chapter tells it that way, as symptoms, the data behind them, and measurements.
Before fine-tuning
In chapter 12, Sprout-1's base model was lost inside the chat template: it had never seen <|user|> or <|assistant|>. Sprout-2's base model is in a different position. In its cooldown, the last fifth of pre-training, 3% of the mix was the fine-tuning conversations written out as plain tokens, 87 million of them, so that the special tokens would not be brand new later; the loss on them fell from 2.59 to 1.59 within 500 steps (chapter 15). The base model has seen the format. It hasn't learned how to behave in it.
To see what fine-tuning changed, we asked every version of Sprout-2 the same 22 questions: 16 simple facts, four skills (who it is, the r's in "strawberry", a shell command, a day of the week) and two sums with big numbers. Each question went in four times at temperature 0.5, first without a system message and then with the one that switches Python on. Start at the base model and step through the versions.
probe.py, with their thoughts and the programs they ran. Each rung is a version of Sprout-2 in the order it was made; the dots count right answers out of four. The table below sums up the whole group of questions.Here is the base model, asked in the chat format, before any fine-tuning:
- "Hi! Who are you?" "I'm [Your Name]. It's nice to meet you. What would you like to talk about?" It knows how an assistant's greeting goes, down to the placeholder in brackets. It has no idea who it is.
- "What is the capital of the USA?" "The capital of the United States is New York City", both times. Overall it gets 70% of the fact questions right, in answers 40 words long on average that wander on after the answer.
- "How many legs does a spider have?" Four in one answer, eight in another: two out of four right.
- "What is 1234 * 5678?" It thinks: "Split 5678 into 500 + 6. 1234 x 500 = 60000…" and ends with 78888. The shape of the thought comes from the conversations it saw in its cooldown; the arithmetic is a guess. The right answer is 7,006,652.
Now switch Python on. The base model writes a program for 94% of the fact questions, and its share of right answers drops to 14%. Asked for the capital of Brazil, it writes a function that sets capital_of_brazil = "Paris", the program runs and prints Paris, and the model reports: "The answer is Paris." In the conversations of its cooldown the tool's system message appeared only where the tool was used, so it learned the rule "tool announced, write a program". A model learns exactly what its data shows, including patterns nobody meant to put there; this chapter will run into that again and again.
A format with more roles
Sprout-1 had four special tokens. Sprout-2's tokenizer keeps the last 16 ids of its 32,768 for special tokens, and nine of them make up the chat format:
| token | id | what it marks |
|---|---|---|
<|endoftext|> | 32752 | the start of every conversation, as of every document |
<|system|> | 32756 | instructions from the page, optional, before everything else |
<|user|> | 32753 | the person's turn |
<|assistant|> | 32754 | Sprout's turn |
<|think|> … <|/think|> | 32757, 32758 | its thoughts, at the start of its turn |
<|tool_call|> | 32759 | instead of an answer: a tool's name on the first line, its input after it |
<|tool_result|> | 32760 | what the tool printed, written by the page |
<|end|> | 32755 | the end of any turn |
The other seven are the three fill-in-the-middle markers from chapter 14, used only in pre-training, and four reserved ids waiting for the agent of part 3. A tool call is plain text after its token: <|tool_call|>python, a newline, and the program exactly as it would be saved to a file. No JSON, no escaped quotes and newlines: a small model writes code more reliably when the code looks the way code always looks.
The loss mask is chapter 12's idea, with two new kinds of tokens inside Sprout's turn. Its thoughts and its tool calls are learned, like its answers and the <|end|> that closes each of its turns. The system message, the question and the tool's output are only read: the page writes the tool result, just as the person writes the question, so there is nothing for the model to learn in predicting it.
sft.py trains on, with the learned tokens underlined. "What the chat sends": the line the page sends at each step of the conversation, and the tokens the model has to write back.The sum with Python, system message included, is 90 tokens, and the loss learns 50 of them; the system message's 12 are only read. With the Python call, the model's turn ends at the program and a second <|assistant|> turn begins after the result; "What the chat sends" shows the same conversation as the two steps the page goes through.
Here is render() from Sprout-2's sft.py, the same function as in chapter 12 with the new roles:
The chat page has its own copy in JavaScript (chatIdsV2 in kit/chatformat.js), and the course's tests check that it produces exactly the ids sft.py trains on: a stray space between how a model was trained and how it is asked, and it is quietly working from a script it never saw. Note what the Python switch in the chat really is. Not a setting inside the model, but this system message, You can run Python code with the python tool., twelve tokens at the start of the line. In the training data a tool call never came without them.
In a training conversation with a tool call, which tokens does the loss teach the model to predict?
The tool result is written by the page, like the question by the person, so the model never has to produce it. Everything the model does write in the chat is learned: its thoughts, its program (that is how it learns to write programs) and its answer, and every <|end|> that tells the page it has finished its turn.
Thinking out loud
Why should a model write its thoughts down at all? Because of how it computes. For every token Sprout-2 writes, all 30 blocks run once: about two operations per parameter, some 250 million, plus attention. That is the same for the word "the" and for the next digit of 1234 × 5678. The model can't stop and think longer about a hard token. What it can do is write: every token it writes goes back in as input, and attention lets every later token read it. A thought of 40 tokens gives the answer 40 more passes through all 30 blocks, and a place to write intermediate results where the next step can find them.
Where do Sprout-2's thoughts come from? Not from an AI assistant writing them for us. There are three sources:
- GSM8K (Cobbe et al., 2021): grade-school word problems with step-by-step solutions written by people.
build_sft.pystrips the calculator notes like<<48/2=24>>and keeps the steps as the thought, the final number as "The answer is 24." - MetaMathQA (Yu et al., 2023): an open dataset of rephrased problems and solutions that its authors produced with GPT-3.5.
- Our own generators: 30,000 problems of 17 kinds from
reasoning.py: column sums, multiplication by splitting a number, primes, counting letters, sorting words, days of the week, clocks, change, discounts, units, reading Python code and more. Code writes the question, the steps and the answer, so they are right by construction.
Here are two of the generators, copied from reasoning.py. Run them, then the cell after them.
"Strawberry" is the first word on the generator's list. Make the 30,000 problems the way build_sft.py does, reasoning.make(30000, seed=1), and 19 of them spell it out, but only one asks about its r's: Count the letter r in "strawberry". In the first widget, pick "strawberry" and step through the versions. From stage 1 on, the thoughts spell the word perfectly, s-t-r-a-w-b-e-r-r-y, and then go wrong: "The letter r is at positions 4, 9. That makes 2." No version got it right more than once in four tries.
The spelling isn't decoration. In the question, " strawberry" is a single token, number 26954, and the model can't see the letters inside it. Spelled out it becomes ten tokens, s -t -r -a -w -b -e -r -r -y, and each r is now there to look at: three copies of one token, -r, number 4504. But the letter in the question is " r", token 377, and to answer the model has to connect the two and then count the matches. Counting identical tokens is known to be hard for attention: it takes a weighted average of what it looks at, and an average over three equal tokens looks just like an average over two.
The day of the week fails the same way: "A week has 7 days, so only the remainder matters: 10 = 1 x 7 + 5. Going forward 5 days from Monday gives Friday." The form is the generator's; the remainder is wrong (10 = 1 × 7 + 3), and five days from Monday isn't Friday either. No version ever found Thursday. And the sums: reasoning.py multiplies numbers up to 99 × 99 and adds numbers up to 9,999, so for 1234 × 5678 the model applies a familiar form to numbers it never saw ("Split 5678 into 567 + 78", which isn't even a split of 5678), and its column addition of 98765 + 43210 misreads the columns and arrives at 141,777.
A thought is a format the model imitates. It gives the model more steps and a place to write, but every step is still one guess of the next token, and it is only as right as those guesses.
How much do thoughts help overall? On GSM8K's 1,319 test problems the chat that thinks gets 10.7% right. The same model writing a Python program gets 25.6%. For arithmetic, there is something better than thinking harder: handing the computation to something that can't get it wrong.
A tool that really runs
Here is the loop. When Sprout-2 decides to use Python, its turn becomes <|tool_call|>python, a newline, a program, and <|end|>. The page notices the <|tool_call|> token, takes the program and runs it. Whatever the program printed goes back as a new turn, <|tool_result|>7006652<|end|>, followed by an open <|assistant|>, and the model goes on writing from there, now with the result in front of it. It may call Python again, up to three times per reply; then it must answer.
The Brazil program runs perfectly and prints Paris. The tool computes exactly what it is given; whether it was given the right thing is still the model's job. Python makes arithmetic exact, not questions sensible. And "a loop that never ends" shows why the page needs a time limit: a program the model writes is a program like any other, and it may never stop.
How the page runs a model's program
The chat uses Pyodide, CPython compiled to WebAssembly, the same Python that runs this course's code cells, but in a separate Web Worker (kit/pytool.js). A program that never ends can't freeze the page: after 8 seconds the worker is thrown away and a new one starts next time. Every call gets a fresh namespace, so one program can't leave variables for the next, and the output is cut to 1,500 characters before it goes back to the model. Everything runs on your device: the page has no server that executes code.
Where did the model learn to call a tool? From TinyGSM (Liu et al., 2023), an open dataset of grade-school word problems solved by short Python programs, all in one style: a function simple_math_problem() with the question in its docstring, a few named variables, a return. We didn't take the programs on trust: build_sft.py ran every program it picked, in a fresh process with a three-second limit, kept only those that finished and printed a number, and used that number as the tool result. In every training example, the tool result is what the program really printed. Stage 2 added arithmetic with big numbers made by code: "Big numbers: I will let Python compute this.", print(1234 * 5678), the result, the answer.
Here is the whole measurement, GSM8K three ways for every fully evaluated version: as the start of a TinyGSM function to finish (the base model's own format), as a chat question with Python off, and with Python on, where the programs really ran.
In its head, the chat gets 11–13% whatever the version; with Python, twice as many. One row stands out: stage 2's first try, where the tool score falls to 17% and the model called Python only 559 times on 1,319 problems. That is the subject of the next section.
Fine-tuning as a series of fixes
First, the data, and who wrote which part of it. It matters, because a small model copies its data very closely:
| data | who wrote it |
|---|---|
| smol-smoltalk (HuggingFaceTB, Apache-2.0) everyday chats, rewriting, summaries, code tasks | larger open models: the dataset is itself synthetic |
| Sprout-1's conversations and 320 new ones about who Sprout-2 is chats at Sprout's level (minus those about Sprout-1's size) | an AI assistant, Claude, to our specification |
GSM8K, MetaMathQA, reasoning.pythoughts before answers | people; MetaMathQA's authors with GPT-3.5; our code |
| TinyGSM, arithmetic Python tool calls | TinyGSM's programs, each run by us; arithmetic made by code |
| 2,850 dialogues (stage 1), 720 writing requests (stage 4) computers, the command line, Python, web data, concepts, geography; emails, letters, notes | an AI assistant, Claude, to our specification |
| fact tables (stage 4) capitals, currencies, chemical symbols, years | mledoze/countries (ODbL), periodictable (NIST data), a list of 60 years in facts.py; questions made by code |
Stage 1: it talks, and invents
The first round is the big one: 258,124 conversations from all of the above except the arithmetic (it came in stage 2) and the stage-4 rows, two passes over them, 1,597 steps, 64 million learned tokens. The held-out loss fell on every kind of data: everyday chats from 1.440 to 1.273, thoughts from 0.354 to 0.246, our new dialogues from 2.598 to 2.228.
Step the first widget to stage 1. Facts are now 84% right, in answers half as long, but with details nobody asked for; a 125-million-parameter model knows the name of a capital far better than the details around it: "The capital of the United States is Washington, D.C. It sits on the Pasco River, just south of the city of Seattle." Switch Python on and the model forgets who it is ("I'm Rohan, a software engineer…") and writes a program for a third of the fact questions: the chemical symbol for gold is gold_symbol = "g". In the browser, in 4 bits, it wrote a program for every question. The cause is the same as with the base model: in stage 1's data the tool message appeared only in tool conversations, and never next to the conversations about who Sprout is.
Stage 2: the tool becomes optional
Stage 2 is short, 266 steps from stage 1's model on 16,990 conversations. 6,320 of them are ordinary chats with the tool message added: the tool is on, and Sprout just answers. 4,766 are cases where calling Python is right: arithmetic with big numbers and TinyGSM problems. 5,904 are a replay of stage 1, so that nothing else is forgotten. The loss on held-out tool arithmetic fell from 2.296 to 0.023 in the first 100 steps.
In the probe it worked: with Python on, the model no longer called it for facts or for who it is, always called it for the sums, and was Sprout-2 again. Then GSM8K came back with 17% instead of 25%, and only 559 tool calls on 1,319 problems. The reason was in the data: the "ordinary chats with the tool on" included 1,718 word problems from MetaMathQA and GSM8K solved in the head, next to 1,916 solved with Python. The model learned exactly that, to use Python on about half of the word problems, and in its head it gets half as many right.
The fix is one line of build_sft.py with a comment that says it all:
With Python on, a word problem with numbers now always goes through Python, 3,416 TinyGSM calls in the redone set; thinking in the head stays for Python off and for puzzles without sums. One more side effect shows in the widget with Python off: from stage 2 on, the model sometimes writes a tool call for the big sums anyway, with nobody to answer it. The chat at the end of the page tells you when that happens and offers to switch Python on.
Stage 2's first try solved 17% of GSM8K with Python on, against 25% for stage 1. What went wrong?
The numbers add up if the programs it did write were as good as before: it just wrote them for fewer than half of the problems (559 calls on 1,319) and did the rest in its head, where it gets about 11% right. Take those conversations out of the "tool on" part of the data and the calls, and the score, come back.
Stage 3: short look-ups, and brevity that spread
The invented details had to go: a small model that adds "one more fact" mostly adds an invented one. Stage 3 cut the answers to the look-up questions of our geography dialogues down to their first sentence ("capital of mongolia" gets "It's Ulaanbaatar."), left the "why", "how" and "what's the difference" questions whole, and showed the 409 short ones twice, 818 conversations among 4,138. We also lowered the chat's temperature from 0.7 to 0.5.
The facts got shorter, from 24 words to 10, and more of them were right. But brevity spread beyond look-ups. "Explain how a rainbow forms" got 13 words; "Tell me about the Moon", between 2 and 34. Only letters and stories kept their length. The second try showed the look-ups once, next to 1,500 long answers to "explain", "tell me about" and "write" requests from openhermes, smol-constraints and our dialogues, so that "answer look-ups briefly" doesn't become "answer everything briefly". Stage 2 redone with this stage 3 on top is the "stages 2–3 fixed" rung: GSM8K with Python back to 23.7% with 1,314 calls, Python never called on facts or on who it is and always on the sums, facts 77–81% right in about ten words.
Stage 4: writing on request, and fact tables
Then a casual request tripped it: "write me an email to my boss asking for a day off" got "I can't accept emails…". The model mixed up writing a text with sending it: Sprout-1's conversations say honestly that it can't send anything. Stage 4 showed 720 new dialogues of casual writing requests twice (emails, letters, notes, poems, toasts; 42 of them "send it for me", answered "I can't send emails, but here is one you can send"), together with the fact tables of the next section and a replay of 4,000 conversations from stage 3: 82 steps.
The email request now gets a full email of 60–110 words, and "write me short letter about llm" 100–170 words instead of 34–76. Fixes have prices here too: placeholders spread ("I'm writing to tell you about [a project or task]"), and explanations got a little worse, their held-out loss up from 1.287 to 1.330. A gentler variant (the writing set once, half the learning rate) was worse on both counts and was not used. GSM8K with Python reached 25.6%, the best of all versions.
Data teaches exactly what is in it. Every fix here is a change of data, and every symptom was in the data before it was in the model.
Facts from a table
Can fine-tuning teach a model facts it doesn't know? Stage 4 tried. facts.py turned tables into questions: the capitals and currencies of UN members, the chemical symbols of the elements, and the years of 60 well-known events, each in three to five phrasings, 1,734 questions in all. A fifth of the countries, elements and events were held out entirely. Afterwards both the trained and the held-out facts were asked in phrasings the training never used ("Tell me the capital of …"), greedily, before and after stage 4.
facts.py saved them, and the case of Brazil.The facts that were in the table improved: capitals from 29% to 46%, currencies from 13% to 23%, chemical symbols from 10% to 44%, years from 40% to 73%. The held-out ones barely moved: capitals 26% to 29%, currencies 11% to 17%, years 25% both times. Symbols are the exception, 26% to 43% even for elements the table never mentioned: the model knew many of them already and learned the answer's form, "The chemical symbol for X is Y", instead of "Gold is a metal with a symbol of eternal life". Facts are learned one by one, and only partly, from three to five examples each.
In the examples, Saudi Arabia and Algeria, from the table, now get their capitals, while Malta and Palau, also from the table, get "Malta" and "Palau". North Korea, held out, moved from Pyongyang to Seoul: fine-tuning shifts facts it never touched. And Brazil was in the table in exactly the probe's words, "What is the capital of Brazil?", with the answer Brasília. The model still says Rio de Janeiro, now in all four tries. A belief built from nine billion tokens of pre-training doesn't yield to five examples.
So should a small model memorise facts at all? A table of 155 capitals, shown five times each, taught it about half of them. The world holds far more facts than 125 million parameters, and each one costs many exposures. The cheaper path is the one this chapter took for arithmetic: don't know it, look it up. A model that can call Python can also call a search, a table or a file. That is the agent of part 3.
The limits of a 125M chat
Three practical limits decided how the chat at the end of this page works.
A long conversation hurts. Sprout-2 learned from short conversations. With 0–2 earlier exchanges in the same conversation it answers 70% of questions right; with 4–8, 56–62%. In one long mixed conversation it drifted to "Washington D.C., on the fourth of the four Atlantic Coasts", while the same question in a fresh chat got the clean answer. So the page shows you the whole conversation but sends the model only the last two questions, with the answer and any tool results in between.
4 bits hurt the chat much more than the base model. The base model paid 0.026 nats for 4 bits per weight (chapter 15); the chat, on the answers of 300 held-out conversations, paid 0.13, five times more. Keeping the embeddings in 8 bits and choosing better clipping won back less than half of it, and no single part of the model was to blame. Worse, the 4-bit chat broke tool calls. 8 bits cost nothing measurable, so the chat ships in 8 bits.
A lower temperature. The chat samples at 0.5 instead of Sprout-1's 0.7, as the probe did: at 125 million parameters, a less likely token is more often an invented detail than a better word.
A word about the system message
A system message is not a command line to the model. It is a dozen tokens at the start of the line, read like any others, and what they do depends only on what followed such tokens in training. Sprout-2 saw a few dozen fixed ones: three ways of announcing the Python tool, "You are Sprout, a small and friendly assistant.", and the task instructions some smol-smoltalk chats carry ("Rewrite the input text…"). It obeys the tool message because thousands of examples showed what follows it, and for the same reason it learned the wrong things from it in stages 1 and 2. A message it has never seen, say "Talk like a pirate", has nothing behind it: chapter 13's pirate was a change of weights, not an instruction. Big chat models follow new system messages because their fine-tuning showed them a huge variety of instructions being followed. A proper chapter on system prompts, with an agent's tools, rules and workspace, belongs to part 3.
Sprout right now
This is the chat from stage 4, the top rung of the first widget. It answers from our server, which runs the same engine and the same 8-bit weights your browser would, so there is nothing to download and your messages are not stored. If the server is busy or out of reach, the chat offers to run Sprout-2 right here: 126 MB, or on a phone the lite 4-bit version, 78 MB, which is noticeably less sharp and may break a tool call. The "Python" switch in the bar is this chapter's system message: on, Sprout-2 is told it can run Python, and its programs run right here in your browser. Thoughts fold away under each answer; the chat remembers the last two questions. Ask who it is, ask for 1234 × 5678 with Python on and then off, watch it spell "strawberry" perfectly and miscount. It knows the capital of Australia and still puts Brazil's in Rio. The last chapter steps back from Sprout: how models like it grow into the ones you use every day, and where to go from here.