4. Training: pretraining, SFT and LoRA

The chapter 3 model continues text. Ask it a question and it continues the question. The step from that to a model that answers is called supervised fine-tuning, and it is small in code and large in effect: the same training loop, a different dataset, and a mask on the loss. This chapter does it twice on the chapter 3 model, once by updating every weight and once by training a low-rank adapter while the weights stay frozen, and measures three things for each: whether the model learned the new skill, how long it took, and what it did to the skill the model already had.

Two of the results are not what the usual account of fine-tuning would lead you to expect, and the chapter is built around them. The numbers come from 04-training.ipynb as executed: 821 seconds in total on the laptop described in chapter 1.

4.1 What pretraining gave us, and what it did not

Chapter 3's model has a held-out loss of 4.807 on the book and knows, in its weights, something about English word order and a great deal about one author. Asked Instruction: reverse the words / Words: tried with from came / Answer:, it produces nothing useful: 0% on every task below. That is not a failure. Nothing in a book trains a model to treat a line starting with Instruction: as a command, and a model only does what its training data made cheap.

Pretraining buys representations. Instruction tuning buys behaviour. The two are different purchases, and the second one is much smaller, which is why a few minutes on a CPU can do it here, and why one study instruction-tuned a 65-billion-parameter model on only 1,000 carefully chosen examples (Zhou et al., 2023).

4.2 The data

Three instructions over short lists of words: reverse the words, repeat the words, and last (output the last word). The answer to each is computed, so a model's output can be checked exactly. The words are the 216 of three to eight letters in the book that the chapter 3 tokenizer encodes as one token when preceded by a space, so each list position is one token; section 4.6 explains why that choice matters. Lists have three to five words, drawn at random.

What is held out is the list. Test lists come from a different random stream than training lists; with 216 words and three to five positions there are over 470 billion possible lists, so a test list has almost no chance of having been seen. The words themselves all appear in training, and section 4.6 shows what happens if they do not.

4.3 The loss mask

A training example is the prompt followed by the answer, as one sequence. The model is trained to predict every answer token, and not trained to predict the prompt. In the batch builder, every prompt position gets a target of -100, which PyTorch's cross-entropy ignores, and the answer positions get the answer tokens.

y = [-100] * (len(p) - 1) + a + [-100]      # predict answer tokens only

Without the mask, the model spends its gradient learning to predict the instruction and the list, which are random by construction and teach nothing. Exercise 4 removes it and measures the cost.

4.4 LoRA

Full fine-tuning updates all 924,416 parameters. A low-rank adapter (Hu et al., 2022) leaves a weight matrix $W$ frozen and adds a trainable correction,

with $r$ much smaller than either dimension. $B$ starts at zero, so the model begins exactly where it was. Here $r$ is 8, $\alpha$ is 16, and the adapters sit on the attention projections qkv and out of every block: 24,576 trainable parameters, 2.7% of the model.

def forward(self, x):
    return self.linear(x) + (x @ self.A.T @ self.B.T) * self.scale

The appeal is obvious: a 2.7% update that can be stored separately, swapped in and out, and trained with a fraction of the optimiser memory. The question is whether 2.7% is enough for a given skill, and this chapter gives a clear answer for this one.

4.5 Training, and the schedule chapter 1 promised

Two thousand steps of 32 examples each, AdamW, for both methods. The learning rate warms up linearly over the first 100 steps and then follows a cosine down to a tenth of its peak:

Warm-up protects the pretrained weights from the first few large, badly estimated updates; the decay lets the model settle. The full model peaks at $10^{-3}$, the adapter at $5 \times 10^{-3}$, higher because it starts from zero and has far fewer parameters to move. Figure 4.1 shows the loss on answer tokens for both runs.

Figure 4.1. Loss on answer tokens, 10-step mean, in nats. Over the last 20 steps full fine-tuning averages 0.236; the adapter is still at 3.773 at the end.

Accuracy is exact match on 100 held-out lists per task, greedy decoding. The two surprises are in the last four columns.

LoRA did not learn the task. At rank 8 on the attention projections, the adapter's loss was still 3.773 at the end of training, and the model it produced reverses nothing. last at 43% is the one task a weak model can partly fake, since the answer is a single token that was in the prompt. This is not the usual LoRA story, and it is worth being precise about why. Reversing a list needs the attention heads to do something they have never done: route each position's query to the right earlier key by position, then carry that token's identity through the value path to the output. That is a new circuit, not an adjustment to an existing one, and a rank-8 change to the projections of a four-block model is too small a change to build it in 2,000 steps. LoRA is strong when the skill is nearly there and the adaptation is a nudge, which is the case for most fine-tuning of large pretrained models, and weak when the skill has to be constructed. Exercise 2 asks at what rank it crosses over, and whether it does.

Full fine-tuning destroyed the model's original skill. The held-out loss on the book went from 4.807 to 30.275, which is far worse than guessing uniformly (6.93). After 2,000 steps on instructions and nothing else, the model that was a passable continuer of Carroll is no longer one. This is catastrophic forgetting, and the number is as bad as it is because the fine-tuning data has no book in it at all. The adapter forgot less, 14.648, because it could not move as far, which is the flip side of its failure to learn. Exercise 5 mixes book text back into the batches, an established way to keep the base skills; InstructGPT, for example, mixes pretraining updates into its reinforcement learning stage to limit regressions on public benchmarks (Ouyang et al., 2022).

Neither of those results would be visible from the training loss, which only reports the task being trained. Measuring the thing you are not training is the point of the last column. Figure 4.2 puts the gains and the costs side by side, with the held-out-word experiment of section 4.6 on the right.

Figure 4.2. Left: exact match on 100 held-out lists per task before fine-tuning (all zero), after full fine-tuning and after LoRA. Middle: loss on the held-out tail of the book, in nats; the yellow line is a uniform guess over 1,024 tokens. Right: the model of section 4.6, on lists of the 156 words it was trained on and on lists of the 60 words it never saw in a training list.

4.6 Holding out words instead of lists

One more run, full fine-tuning again, but with 60 of the 216 words removed from every training list and the test lists built only from those 60.

The training loss reached 0.129, against the first run's 0.236. The model learned to reverse, as the first row and Figure 4.2 (right) show. It cannot reverse a word it never saw in a training list, as the second row shows, even though the word is in its vocabulary and the model saw it many times while it was trained on the book. Copying a token through attention is a learned path through the value and output projections, and that path was trained for 156 tokens and not for the other 60. "Generalisation" always has a unit, and here the unit the model generalises over is the list, not the word. This is the same shape as chapter 1's result on longer strings: what a model has learned is bounded by what it was shown, along every axis you did not vary.

Where this goes wrong

Training on the prompt. Leaving the mask out costs capacity and, worse, makes the loss look better than it is, since the fixed parts of the prompt (Instruction:, the words, Answer:) are trivially predictable and pull the average down. Compare losses only with the same mask.

Reading the training loss as the result. The adapter's loss curve goes down. The full model's goes down further. Neither says anything about accuracy on held-out lists or about the book. Every row of the table came from a separate measurement.

Assuming the adapter is a cheaper version of the same thing. It is a cheaper version of a different thing: a constrained update. On this task the constraint was binding. The only way to know is to run both.

Exercises

1. Put the adapters on the MLP layers instead of the attention projections, same rank. Report trainable parameters and accuracy.
2. Change the rank to 2, 32 and 64. At what rank, if any, does LoRA match full fine-tuning on this task, and what does that cost in parameters?
3. Train on reverse and repeat only, then test last. Report accuracy on the untrained instruction. What does the model do instead?
4. Remove the loss mask so that the prompt tokens are also trained on. Report accuracy and the held-out book loss. Which number moved, and why?
5. The full model's held-out book loss rose after fine-tuning. Mix one book sequence into every batch of instruction examples and retrain. Does the book loss recover, and does accuracy pay for it?

Notebook for this chapter: 04-training.ipynb, 821 seconds in total, no downloads. It writes checkpoints/04-sft.pt, which chapters 5, 6 and 8 use.

Further reading

• Edward J. Hu, Yelong Shen, Phillip Wallis, et al. 2022. LoRA: Low-rank adaptation of large language models. International Conference on Learning Representations (ICLR 2022). arXiv:2106.09685.
• Chunting Zhou, Pengfei Liu, Puxin Xu, et al. 2023. LIMA: Less is more for alignment. Advances in Neural Information Processing Systems 36 (NeurIPS 2023), 55006-55021. arXiv:2305.11206.
• Long Ouyang, Jeff Wu, Xu Jiang, et al. 2022. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems 35 (NeurIPS 2022), 27730-27744. arXiv:2203.02155.