Teaching it to talk

Ask base Sprout “Why is the sky blue?” It may respond with another question, begin a story or continue a scene between two characters. It sometimes answers too: its corpus included dialogues. But it was trained to predict document continuations, not to consistently act as your conversation partner. Our next task is to make an appropriate assistant reply more likely to follow a user's question.

We can do this without changing the architecture. We start from the trained weights and continue training on prepared conversations. Alongside the data, we change how loss is counted: the assistant's answers become the targets. This stage is called supervised fine-tuning, or SFT, a common way to prepare a language model for chat. Sprout has about twenty thousand conversation records for it, including repeated examples.

A storyteller doesn't answer

Below, the same question goes to two versions of Sprout: the base model from chapter 10 and the same model after fine-tuning on conversations. The base model can be asked in three different ways.

Try all three ways of asking the base model.

• Just the question. The base model treats it as the opening of a text and carries on: another line of dialogue, a narrator, sometimes a story. Nothing in its training said that a question at the start of a document must be followed by its answer.
• As a script. Here we wrap the question in a tiny scene: "Anna talks to her friend Sprout. Anna: … Sprout:". Sprout read plenty of such scripts in the SODA dialogues (chapter 9), and in a script the line after "Sprout:" is naturally a reply. Often this works surprisingly well. This trick, dressing the task up as a document whose natural continuation is the answer, was the main way to use language models before chat fine-tuning. GPT-3 (Brown et al., 2020) was used almost entirely this way.
• In the chat template. The question is wrapped in the same special marks the chat model gets. The base model has never seen them, so it is lost: it often starts in the middle of a word ("es can be bright…", "holders, I'm Sophia"), then drifts into a story or a dialogue between strangers that has little to do with the question.

The key observation is on the other side: the fine-tuned model answers. Yet it has exactly the same architecture, the same 17.31 million parameters, and was trained on only a small amount of extra text. The base model already knew grammar, how dialogues go, and the simple facts of a children's world. What it lacked was a format and a role: knowing that it is one side of a conversation and must answer the other side.

A script with special marks

The "script" trick hints at the solution: turn a conversation into a document of a fixed shape, and train the model on many such documents. We just need a shape that can't be confused with ordinary text. "Anna:" and "Sprout:" could appear in any story; we need marks that never occur in normal writing.

That is what special tokens are for. Back in chapter 6, when we built the tokenizer, we set aside four ids at the very end of the vocabulary: 8188 to 8191. Only one of them, <|endoftext|>, has been used so far, to separate documents in pre-training. Now the other three come into play:

• <|user|>: the person's turn starts here;
• <|assistant|>: Sprout's turn starts here;
• <|end|>: the current turn is over.

A whole conversation becomes one line of tokens:

<|endoftext|><|user|>hi!<|end|><|assistant|>Hello! How is your day going?<|end|><|user|>…

To chat, we build this line from the conversation so far, finish it with <|assistant|>, and let the model write until it produces <|end|>, the stop token promised in the previous chapter. Every chat box on this site does exactly that. Build a conversation of your own below and see how it turns into tokens:

What happens in a long conversation?

In our chat, sending a message rebuilds the available history as a token sequence. That sequence carries the conversation; chatting does not update the model's weights. Sprout has a 512-token context, so the interface keeps recent turns within a 440-token budget before starting a reply. Older text is dropped, explaining forgotten names or agreements in a long conversation. Other systems have different windows and history-management strategies: they may summarise earlier exchanges or retrieve information from a separate store.

This is the function from sft.py that turns a conversation into training material. Besides the token ids it builds a second list of the same length, the mask: 1 where the model must learn to predict the token, 0 where it only reads.

def render(tok, messages): """Token ids and a 0/1 mask saying which positions the model should learn to predict.""" ids, mask = [tok['<|endoftext|>']], [0] for m in messages: role = tok['<|user|>'] if m['role'] == 'user' else tok['<|assistant|>'] body = tok.encode(m['content'].strip(), allow_special=False) + [tok['<|end|>']] ids += [role] + body learn = 1 if m['role'] == 'assistant' else 0 mask += [0] + [learn] * len(body) return ids, mask

The allow_special=False setting prevents ordinary message text from becoming structural markup. If a user types “<|end|>”, the tokenizer encodes ordinary text; the program inserts the actual marker id itself. This separates message content from role and boundary markers. It prevents marker spoofing through tokenization, but does not solve prompt injection in general: unwanted instructions can also be written in ordinary words.

Grading only the answers

In pre-training every token was a target: the loss was the average surprise over all positions. In fine-tuning we multiply each position's surprise by the mask and average only over the positions that count:

$$\mathcal{L} \;=\; \frac{\sum_t m_t \cdot \big(-\ln p_\theta(x_t \mid x_{\lt t})\big)}{\sum_t m_t}, \qquad m_t \in \{0, 1\}.$$

The model still reads every token. The user's words pass through the network and affect the reply. The mask removes the direct loss for predicting user tokens, but gradients from the answer can still flow through computations involving the question. Thus “train on the answers” describes the loss, not a ban on learning to process the input. Why choose this mask?

1. The target is the assistant reply. The mask directs the loss towards the part of the conversation that the model will generate in use.
2. The question is already given. At generation time it is an input condition, not text the model must write on the user’s behalf. Predicting its tokens separately is therefore optional for our task.
3. Control over weighting. Long questions do not gain more weight in the loss simply by being longer. Training without this mask is also possible; the better choice depends on the data and task.

The special tokens play different roles. The program inserts <|assistant|>, so our loss does not score its prediction. The <|end|> closing an answer is scored along with the answer, teaching the model when to emit a completion signal. Excluding it would remove this direct training signal for stopping at the right point. An external length limit could still interrupt generation.

Here is the whole idea on a toy example. The probabilities are made up; watch how the user's word "sprout", which the model did not expect at all, drags the plain average up, and how the mask removes it.

import numpy as np # a toy conversation, one "token" per word; the mask marks what the model is graded on tokens = ["<|endoftext|>", "<|user|>", "hi", "sprout", "<|end|>", "<|assistant|>", "hello!", "how", "are", "you?", "<|end|>"] mask = np.array([0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1]) # the probability the model gave to each actual token (made up for the example) p = np.array([1.0, 0.9, 0.02, 0.001, 0.3, 0.95, 0.4, 0.6, 0.9, 0.8, 0.7]) surprise = -np.log(p) print("loss over every token: ", round(surprise[1:].mean(), 2)) print("loss over Sprout's turn:", round((surprise * mask).sum() / mask.sum(), 2)) print() # the first token is never predicted, so start from the second for tok, s, m in zip(tokens[1:], surprise[1:], mask[1:]): print(f"{tok:>14} {s:6.2f} {'graded' if m else '-'}")

In Sprout's model.py the same thing takes two lines at the end of forward: the cross-entropy is computed for every position with reduction='none', multiplied by the mask and divided by the number of ones.

Twenty thousand conversations

The template and the mask are the mechanism. What the model becomes depends on the conversations we feed it. Sprout's collection has 19,734 of them from four sources:

• 7,217 conversations written for Sprout's level across 18 topics: small talk, nature, science, daily life, stories, learning, friends, Sprout itself, the world, conversations that depend on earlier turns, technology, history, how-to, fun, health, the planet, simple explanations and a mixed bag. They follow one style guide: short sentences, words a ten-year-old knows, plain text without lists or emoji, and honesty about limits. Sprout says that it is small, that it only understands English, that it doesn't know today's date or the news, and that it can make mistakes. Each conversation appears twice in the mix, and the 400 about Sprout itself four times, so the model firmly learns who it is.
• 1,200 everyday conversations from the open everyday-conversations dataset: longer back-and-forth chats with an assistant.
• 2,500 story requests. The request names the hero of a TinyStories story ("Tell me a story about a little fox"), and the answer is that story, at most 170 words.
• 800 "continue this story" requests: the first sentence or two of a TinyStories story in the request, the rest in the answer.

Switch the strip to "by Sprout's tokens" and the picture changes. Written conversations are 77% of all conversations but only half of the graded tokens: they are short, 64 tokens on average, while a story request averages 178. Since the loss is an average over graded tokens, each source pulls on the model in proportion to its share of graded tokens, not of conversations.

That arithmetic caught us out. Our first attempt used 5,000 stories, up to 230 words long. The fine-tuned model was fluent and polite, but when someone wrote "I feel a bit sad today", it answered with a fairy tale. Of course: stories made up the biggest share of everything it had been graded on, so "answer with a story" had become its default. We halved the number of stories and kept only the short ones. Then came a second, subtler trap. The first story requests were built from SimpleStories' labels such as "bygone eras" or "an anti-hero", and the stories only loosely matched them. Asked for "a story about a brave little fox", the model cheerfully told one about a girl baking a cake: it had learned that a story need not fit the request. We rebuilt the requests from each story's own hero ("Once upon a time, there was a little fox…" becomes "Tell me a story about a little fox"), and the fox came back. (Ask "I feel a bit sad today" in the first widget to see how the final model takes it.) Data mixing is one of the least glamorous and most important decisions in building a language model.

Here is how the collection is assembled. conversation() checks that a chat starts with the user, ends with the assistant and alternates properly; story_requests() finds each story's hero in its first sentence and asks for exactly that story:

def conversation(messages, source): msgs = [{'role': m['role'], 'content': tidy(m['content'])} for m in messages if m['role'] in ('user', 'assistant')] ok = msgs and msgs[0]['role'] == 'user' and msgs[-1]['role'] == 'assistant' and all(m['content'] for m in msgs) ok = ok and all(a['role'] != b['role'] for a, b in zip(msgs, msgs[1:])) return {'messages': msgs, 'source': source} if ok else None STORY_ASKS = [ 'Tell me a story about {subject}.', 'Can you write a story about {subject}?', 'story about {subject} pls', 'I want a short story about {subject}.', 'Tell me a bedtime story about {subject}.', 'Write a story where the hero is {subject}.', 'Could you tell me a story about {subject}, please?', 'Make up a story about {subject}!', ] # "Once upon a time, there was a little fox named ..." -> "a little fox" SUBJECT = re.compile(r"\bthere (?:was|lived) (a|an) ([a-z]+(?: [a-z]+){0,3}?)(?= named| who| that| in | with| called|,|\.)") def story_requests(raw, n, rng, per_noun=120): """Ask for a story about exactly the hero the story is about, so request and answer agree. (A first version asked for SimpleStories' abstract topic labels like "bygone eras"; the stories only loosely matched them, and the model learned that a story need not fit the request.) """ with open(f'{raw}/TinyStoriesV2-GPT4-train.txt', encoding='utf-8') as f: text = f.read(150_000_000) per = {} out = [] stories = text.split('<|endoftext|>')[1:-1] rng.shuffle(stories) for s in stories: s = tidy(s) m = SUBJECT.search(s[:220]) if not m or len(s.split()) > 170 or not 1 <= len(m.group(2).split()) <= 3: continue noun = m.group(2).split()[-1] if per.get(noun, 0) >= per_noun: continue per[noun] = per.get(noun, 0) + 1 ask = rng.choice(STORY_ASKS).format(subject=f'{m.group(1)} {m.group(2)}') out.append(conversation([{'role': 'user', 'content': ask}, {'role': 'assistant', 'content': s}], 'story')) if len(out) >= n: break return out

Altogether, the conversations come to 1.8 million tokens, and 75% of them belong to Sprout's turns. Compare that with the 330 million tokens of pre-training: the fine-tuning data is about half a percent of what the model has already read.

Fine-tuning itself

sft.py is a close cousin of train.py from chapter 10. It starts from the base model's weights instead of random ones; otherwise it is the same model, the same optimisers (Muon for the matrices, AdamW for the rest) and the same forward pass, now with the mask. What differs is the scale and the care:

• The learning rate is five times lower than in pre-training (these are the defaults in sft.py): 0.004 instead of 0.02 for the matrices, 0.0008 instead of 0.004 for the embeddings. The model already knows a lot, and big steps would trample it.
• By default the data is shown three times (three epochs) in batches of 32 rows, with 20 warm-up steps and a linear fade of the learning rate over the second half.
• We hold out 3% of the conversations. They are never trained on, only used to measure the loss, so we can see whether the model is learning to talk or just memorising.

Conversations have different lengths, and the model trains on rows of fixed length. pack() glues conversations one after another into rows of exactly 513 tokens: 512 inputs plus one extra, because the targets are the same row shifted by one position. Whatever doesn't fit starts a new row, and the leftover space is filled with padding, which the mask ignores.

def pack(examples, context): """Glue conversations into rows of exactly context+1 tokens (the rest is padding).""" rows, cur_ids, cur_mask = [], [], [] for ids, mask in examples: if len(ids) > context + 1: continue if len(cur_ids) + len(ids) > context + 1: rows.append((cur_ids, cur_mask)) cur_ids, cur_mask = [], [] cur_ids += ids cur_mask += mask if cur_ids: rows.append((cur_ids, cur_mask)) x = torch.zeros(len(rows), context + 1, dtype=torch.long) m = torch.zeros(len(rows), context + 1, dtype=torch.long) for i, (ids, mask) in enumerate(rows): x[i, :len(ids)] = torch.tensor(ids) m[i, :len(mask)] = torch.tensor(mask) return x, m

Here is the same packing done in your browser on the sample conversations from the explorer above:

Sprout's conversations are short: the median is 66 tokens and the longest is 359, so none has to be thrown away. The training part packs into about 3,900 rows that are nearly 90% full; with the default settings that makes about 370 steps. Next to the 10,070 steps of pre-training, that is a light touch.

How many epochs? Too few and the new format doesn't stick; too many and the model starts memorising particular conversations instead of learning how to converse. The held-out conversations act as referee. While their loss goes down, the model is generalising. If it turns up while the training loss keeps falling, the model is memorising. Here is Sprout's real fine-tuning run:

What changed inside

The loss curve is an average. To see the change up close, let's give both models the same conversation, one they have never seen, and measure how surprised each is by every token of Sprout's reply. This is exactly the quantity fine-tuning pushes down.

Two things stand out. The ordinary words of the reply become less surprising, especially the first ones: the base model doesn't expect an answer to begin right after <|assistant|>, because it has never seen that token. And the final <|end|> changes most of all. For the base model it is a token that never once appeared in its training data, so its probability is vanishingly small. After fine-tuning, Sprout expects to close its turn at the right moment: in the first two conversations it gives <|end|> more than 99%. The frog is the exception worth a look. After a story of just two sentences the chat model still expects more (the stories it was trained on went on for a paragraph or two), so <|end|> gets only a few chances in a hundred thousand. Even that is a thousand times more than the base model gives it.

Why the base model thinks <|end|> is nearly impossible

For softmax and cross-entropy, the derivative with respect to logit $z_j$ is $p_j-y_j$. If a token is never the correct target, $y_j=0$, and this direct loss signal always favours lowering its score. That is the situation for chat markers during pre-training. But logits are computed from shared weights, not learned as independent parameters, so an actual optimizer step need not lower every such logit. The base model also had no practice reading these markers as input. Fine-tuning provides examples of their conversational use and, for the closing <|end|>, positive prediction targets as well.

The weights themselves moved only a little: a few hundred steps at a low learning rate. But the behaviour changed completely. That is the pattern of fine-tuning in general: a small nudge in weight space, a big change in what the model does.

Why is the <|end|> that closes Sprout's turn graded, while <|assistant|> is not?

<|assistant|> is too rare to learn We insert <|assistant|> ourselves, but only the model can decide when its answer is finished Special tokens are never graded; <|end|> is an exception for technical reasons

When we chat, our code appends <|assistant|> to hand the turn to Sprout, so predicting it would be pointless. The end of the answer is Sprout's own decision, so we include this marker among the targets. An external length limit can also stop generation.

What fine-tuning can't do

What does the model learn in fine-tuning? In our case, answer format, tone and stopping behaviour are especially visible. In LIMA (Zhou et al., 2023), researchers fine-tuned a 65-billion-parameter LLaMA on a thousand carefully chosen examples and obtained strong answer ratings. They proposed that, with a capable base model, much of this tuning teaches it to use existing abilities in the desired format. That is a hypothesis about their setting, not a claim that fine-tuning cannot add knowledge or skills.

Fine-tuning can teach new information, but a small collection of conversations does not turn Sprout into a reliable calculator. Try “What is 17 times 23?” The correct answer is 391; recorded trials instead produced confidently phrased, incorrect numbers. Getting the format right and getting the content right are separate achievements. A hallucination is generally a plausible-sounding output that is false or unsupported by the context. Examples of acknowledging limitations can help, but do not guarantee that a model will recognise every error. Gudibande and colleagues (2023) likewise found that, in their experiments, imitating a larger model transferred its style more readily than its factual accuracy. We therefore need to assess Sprout's conversational behaviour and the correctness of its answers separately.

Sprout right now

This is the Sprout we have been growing all along: the same 17.31 million parameters, now fine-tuned to hold a conversation. Say hello, ask it about animals or the weather, ask for a story, tell it how your day went. It is small, so it sometimes gets facts wrong or loses the thread in long conversations, and it only understands English. In the next chapter you will learn to fine-tune it yourself, giving it a style of its own without retraining all 17 million numbers.