Your own fine-tune

Sprout can talk now. It answers questions, tells stories and explains politely that it only understands English. But it talks exactly the way we taught it in chapter 12: plainly, kindly, a little blandly. What if we want Sprout to talk like a pirate? Or to answer in verse? Or to know your family's recipes, which nobody else has ever written down?

The obvious answer is to fine-tune it once more, the same way we taught it to chat. That works. But look at the price: every new manner of speaking becomes another full copy of all 17 million numbers, and training it needs memory for each of those numbers several times over. For Sprout that is bearable. For a model with eight billion parameters it is a server room.

This chapter is about the trick that made fine-tuning available to anyone with a laptop: LoRA, low-rank adaptation. The idea fits in a sentence: leave the old weights untouched and add a thin detour beside each of them. To see why the detour can be thin, we'll spend some time with one of the most beautiful ideas in linear algebra, the rank of a matrix.

Fine-tuning everything

Recall what happened in chapter 12. We took the base Sprout and kept training it, now on conversations: the same loss, the same descent, only new data. Each of the 17,309,056 parameters was free to change. This is called full fine-tuning.

What does it cost? During training the optimiser keeps far more than the weights. Every parameter needs a gradient, and Adam also keeps two running averages per parameter (the moments from chapter 10). With the usual mixed-precision setup this adds up to about 16 bytes per parameter. For Sprout: 17.31 million × 16 bytes ≈ 277 MB, which a laptop won't even notice. For Llama 3 8B the same arithmetic gives about 128 GB, more than even an 80 GB data-centre card holds.

There is a second cost too. The result is a whole new model. Want a pirate and a poet? Two more copies. A hundred customers, each wanting their own tone? A hundred copies.

And yet intuition says that a new style is a small change. Pirate Sprout still knows grammar, still knows the Moon goes around the Earth, still puts words in the right order. Only the manner changes: "arr", "matey", "be" instead of "is". Does a small change really need 17 million new numbers? To answer that, we first need a way to measure how complicated a change to a matrix is.

How many numbers does a change need?

Every weight in Sprout lives in some matrix, a rectangular table of numbers. The attention layer of each block, for example, holds a 384 × 1152 matrix: 442,368 numbers. The easiest way to get a feel for a table of numbers is to see it as a picture, where each number is the darkness of a pixel. So let's play with pictures first.

Take a column of numbers $u$ (one per row) and a row of numbers $v$ (one per column), and fill the table with the rule "cell = its row's number × its column's number":

$$W_{ij} = u_i \, v_j, \qquad W = u\,v^\top .$$

This is an outer product, and the picture it makes is very plain: every row is the same pattern $v$, made darker or lighter by the factor $u_i$. It looks like tartan, with threads crossing at right angles. Such a matrix has rank one, and storing it takes $m + n$ numbers instead of $m \cdot n$.

Add up several such tartans and you can draw almost anything. The smallest number of rank-one layers that add up exactly to a matrix is called its rank, $r$. A matrix of rank $r$ can be stored in $r(m + n)$ numbers: $r$ columns and $r$ rows.

Play for a minute and three things stand out. Sprout's picture becomes recognisable after just a handful of layers, when the right-hand image uses only a small fraction of the numbers. The tartan is restored exactly at rank 3, because that is how it was made: from three outer products. And noise refuses to shrink. Its singular values barely fall, each layer carries about as much as the previous one, and the error melts away only as the rank approaches the full 64. Random numbers have no structure, so there is nothing to compress.

The same thing in numpy: the SVD is one line, the best rank-$r$ copy is another. Run it, and check the last line against the theorem.

import numpy as np m, n = 64, 96 i, j = np.mgrid[0:m, 0:n] # a "picture": a round sun, a striped sky, a dark strip of ground W = 1.0 * ((i - 24) ** 2 + (j - 60) ** 2 < 14 ** 2) W += 0.3 * np.sin(i / 5) * np.cos(j / 9) W[54:, :] = 0.8 U, S, Vt = np.linalg.svd(W, full_matrices=False) for r in [1, 2, 4, 8, 9, 10]: W_r = (U[:, :r] * S[:r]) @ Vt[:r] # the best rank-r copy of W err = np.linalg.norm(W - W_r) / np.linalg.norm(W) print(f"rank {r:2d}: {r * (m + n):5d} numbers instead of {m * n}, error {err:5.1%}") # Eckart-Young: the error is exactly what the dropped singular values carry r = 8 print(f"check: {np.sqrt((S[r:] ** 2).sum()) / np.linalg.norm(W):.1%}")

At rank 10 the error drops to zero, and 1,600 numbers store the whole picture exactly. That's no accident: every row of this picture is a mix of just ten patterns. A disc drawn on a grid has only eight different slice widths, so the sun gives eight patterns; the wave adds one more and the ground another. Rank measures exactly this: how many independent patterns a table really holds.

Now the question that matters. Fine-tuning changes the weights: $\Delta W = W_\text{after} - W_\text{before}$. That change is a matrix too. What is its rank? In 2020 Armen Aghajanyan and colleagues showed that fine-tuning has a surprisingly low "intrinsic dimension": a language model can be tuned for a new task by moving surprisingly few numbers. A year later Edward Hu and colleagues at Microsoft drew the practical conclusion: if the change is simple anyway, why not look for it in low-rank form from the start?

LoRA: a detour beside the weights

Here is all of LoRA. A layer of Sprout computes $y = Wx$. We freeze $W$: not one of its numbers will change. Beside it we put a detour made of two thin matrices:

$$y = W x + \frac{\alpha}{r}\, B A\, x, \qquad A \in \mathbb{R}^{r \times d_\text{in}},\quad B \in \mathbb{R}^{d_\text{out} \times r}.$$

$A$ squeezes the input down to $r$ numbers, and $B$ unfolds them back to the size of the output. Their product $BA$ is a $d_\text{out} \times d_\text{in}$ matrix, the same shape as $W$, but its rank is at most $r$. Only $A$ and $B$ learn. For the 384 × 1152 attention matrix, a rank-8 detour holds $8 \cdot (384 + 1152) = 12{,}288$ numbers instead of 442,368, which is 2.8%.

The factor $\alpha / r$ is a volume knob. Thanks to it, changing the rank doesn't change the typical size of the update, so the learning rate doesn't need to be retuned. In our lora.py $\alpha = 16$ and $r = 8$, so the detour comes in with a factor of 2.

There is one subtle point: how to start. LoRA fills $A$ with small random numbers and sets $B$ to exactly zero. Then $BA = 0$ at the first step, and the model is precisely the chat model: the adapter begins by doing nothing and learns from there. But why not set both to zero? Let's work out the gradients.

You can check all of this, and watch the whole process, on a toy. Sprout's picture from above plays the frozen weight $W$. The "fine-tune" is a costume: the target is Sprout dressed as a pirate or as a poet. The change $\Delta W$ is the costume itself, and a rank-$r$ LoRA has to find it by gradient descent, starting from $B = 0$.

Different costumes cost different ranks. The poet wears a beret and a scarf, mostly horizontal and vertical stripes, and by rank 8 the costume sits almost perfectly. The pirate's eye-patch strap runs diagonally, and a diagonal line is one of the most expensive things for low rank: every row has its own dark dot in its own place, the same effect you saw when drawing. At the same rank 8 the pirate's error is still more than three times the poet's. And notice something else: whatever the rank, descent arrives at the dashed line, the Eckart–Young bound. Two thin matrices trained by plain descent find the best low-rank change on their own.

What would happen if both A and B started at zero?

The adapter would learn twice as fast Nothing: both gradients would be zero, and the adapter would never move The model would forget everything on the very first step

The gradient for A contains B, and the gradient for B contains A. If both are zero, every step of descent is zero. One of the two must start random; LoRA picks A, so that the product BA still starts at zero and the model starts out unchanged.

What it saves

Now let's do the sums for real models. Our lora.py puts a detour beside every matrix inside the blocks: attention's combined Q K V matrix and its output projection, and all three matrices of the feed-forward layer. The embedding table is left alone. Choose a model, a rank and where the adapters go:

For Sprout, LoRA is a convenience rather than a necessity: both variants fit on any laptop. At scale, the picture changes. For Llama 3 8B, full fine-tuning needs well over a hundred gigabytes, while with LoRA the frozen weights take 16 GB and the adapter only a few hundred megabytes on top. With 4-bit frozen weights, the QLoRA trick, the whole thing shrinks to the memory of an ordinary laptop. This is how people fine-tune open models at home.

In code it's even shorter than in words. The whole LoRA layer from lora.py:

class LoRALinear(nn.Module): def __init__(self, base, rank=8, alpha=16): super().__init__() self.base = base self.A = nn.Parameter(torch.randn(rank, base.in_features) / math.sqrt(base.in_features)) self.B = nn.Parameter(torch.zeros(base.out_features, rank)) self.scale = alpha / rank def forward(self, x): return self.base(x) + (x @ self.A.T @ self.B.T) * self.scale def merge(self): self.base.weight.data += (self.B @ self.A).to(self.base.weight.dtype) * self.scale return self.base

Three details match the maths above. A starts random with scale $1/\sqrt{d_\text{in}}$, and B starts at zero. forward multiplies by A.T first and only then by B.T: the full $d_\text{out} \times d_\text{in}$ matrix $BA$ is never built, and the detour costs only $r(d_\text{in} + d_\text{out})$ multiplications per token. And merge is a single line; we'll get to it in a moment.

Wrapping the model is just as short: freeze everything, then replace each target layer with its wrapped version. The optimiser receives only the list of new matrices.

TARGETS = [('attn', 'qkv'), ('attn', 'proj'), ('ffn', 'w1'), ('ffn', 'w2'), ('ffn', 'w3')] def add_lora(model, rank, alpha): """Freeze everything, then wrap every projection in the blocks with a LoRA detour.""" for p in model.parameters(): p.requires_grad_(False) for block in model.blocks: for part, name in TARGETS: parent = getattr(block, part) setattr(parent, name, LoRALinear(getattr(parent, name), rank, alpha)) return [p for p in model.parameters() if p.requires_grad] # ... trainable = add_lora(model, args.rank, args.alpha) # ... opt = torch.optim.AdamW(trainable, lr=args.lr, weight_decay=0.0)

Everything else in lora.py is the training loop from chapter 12: the same render() and pack(), the same loss counted only on Sprout's own words.

Merging it back in

After training we have the frozen $W$ and the detour $BA$. At inference time we could compute both paths for every token. But the detour is linear, so the two paths can be glued into a single matrix:

$$W x + s\,BAx = (W + s\,BA)\,x = W' x .$$

We compute $W' = W + s\,BA$ once, write it where $W$ used to be, and throw the detour away. What remains is an ordinary Sprout: the same shapes, the same 17,309,056 numbers, the same speed. That's why Pirate Sprout in this chapter is a regular 17.5 MB model file, which the browser engine loads without knowing anything about LoRA.

def merge_lora(model): for block in model.blocks: for part, name in TARGETS: parent = getattr(block, part) layer = getattr(parent, name) if isinstance(layer, LoRALinear): setattr(parent, name, layer.merge()) # ... adapter = {n: p.detach().cpu() for n, p in model.named_parameters() if p.requires_grad} torch.save({'rank': args.rank, 'alpha': args.alpha, 'adapter': adapter}, f'{args.out}/adapter.pt') merge_lora(model) torch.save({'config': cfg.to_dict(), 'model': model.state_dict(), 'step': steps}, f'{args.out}/ckpt_final.pt')

Our lora.py saves both: the adapter on its own (adapter.pt, the 417,792 numbers) and the merged model. Merging isn't the only option. If you keep adapters separate, one base model in memory can serve many styles at once, attaching a different small file to each request: S-LoRA (Sheng et al., 2023) serves thousands of adapters from one GPU this way. We merge because Sprout is small and the browser engine should stay simple.

Now let's put the whole idea together in about twenty-five lines of numpy you can run: a frozen layer, a hidden rank-2 change the adapter has to discover, a rank-4 detour trained with exactly the gradients we derived, and the merge at the end.

import numpy as np rng = np.random.default_rng(1) d_in, d_out, r = 64, 48, 4 W = rng.standard_normal((d_out, d_in)) / 8 # frozen, "pre-trained" weight # the new habit: a rank-2 change the adapter has to discover dW = rng.standard_normal((d_out, 2)) @ rng.standard_normal((2, d_in)) / 8 X = rng.standard_normal((256, d_in)) # 256 inputs Y = X @ (W + dW).T # what the tuned layer should say A = rng.standard_normal((r, d_in)) / np.sqrt(d_in) # random, as in lora.py B = np.zeros((d_out, r)) # zero: step 0 changes nothing s, lr = 2.0, 0.01 # s = alpha / r; lora.py has 16 / 8 for step in range(151): h = X @ A.T # the detour: x -> A x Y_hat = X @ W.T + s * h @ B.T # frozen path + s * B A x err = Y_hat - Y loss = (err ** 2).sum() / len(X) g = 2 * err / len(X) # dLoss / dY_hat dB = s * g.T @ h # s * g (Ax)^T dA = s * (g @ B).T @ X # s * B^T g x^T A -= lr * dA B -= lr * dB if step % 25 == 0: print(f"step {step:3d} loss {loss:.5f}") W_merged = W + s * B @ A # merge: an ordinary matrix again print("merged gives the same outputs:", np.allclose(X @ W_merged.T, X @ W.T + s * (X @ A.T) @ B.T)) print("trained", A.size + B.size, "numbers instead of", W.size)

The loss falls from 89 to zero, and after merging the layer gives the same outputs as the frozen path plus the detour. Now change the last number in d_in, d_out, r = 64, 48, 4 to 1 and run it again: the loss gets stuck at about 41. The change we are looking for has rank 2, and a rank-1 detour cannot express it no matter how long it trains.

Sprout's three voices

Time for the real thing. We trained two adapters on top of the chat model, sprout-chat. The data for each is 500 ordinary Sprout conversations in which every answer was rewritten: for one adapter in pirate talk, for the other in rhyme. The questions stayed as they were; only the answers changed their clothes. The settings are lora.py's defaults: rank 8, α = 16, AdamW with a learning rate of 0.003, nine passes over the data, and the loss counted only on Sprout's words, as in chapter 12. We found the number of passes by watching the held-out conversations: their loss bottoms out around the seventh or eighth pass and then climbs again, because 500 conversations are easy to memorise. That is only 34 steps for the pirate and 41 for the poet: less than a minute on the Mac.

Ask all three the same question. They share one random seed, so the only difference between their answers is the weights.

How much did each adapter have to learn? Here are their real training curves. The thick line is the loss on conversations the adapter never trained on (lora.py sets one in twenty aside for this); the thin one is the loss on training batches.

Compare where the curves start and how far they fall. The starting point is the chat model's surprise at the new manner of speaking; the end point shows how much of that surprise rank 8 managed to absorb from a few hundred examples. The gap between the thick and thin lines is a hint of memorisation: if the training loss keeps falling while the loss on unseen conversations stands still, the adapter is learning the particular sentences rather than the style.

What the model forgets

Whatever a network learns, it learns by changing numbers that were already doing something. When fine-tuning pulls the weights towards a new task, old skills can get worse. In the extreme this is called catastrophic forgetting: a network trained on task B forgets task A altogether. McCloskey and Cohen described it back in 1989 on the simplest networks, and it has never gone away.

Forgetting can be measured with the tool from chapter 2: surprise. Let four versions of Sprout read the same texts and compute their average surprise, $-\ln p$ per token. A model that has "forgotten" how to tell ordinary stories will be more surprised by an ordinary story.

What to look for: if the adapters learned their styles, Pirate Sprout should be the calmest model in the pirate row, and Poet Sprout in the verse row. The price is in the other rows: how much more surprised is each adapter by a plain answer or a plain story than the chat model it grew from? That difference is forgetting, measured in nats. On our models both adapters did learn their styles (2.93 against the chat model's 4.14 in the pirate row, 3.72 against 5.14 in the verse row), and the forgetting is small but uneven: the pirate is only 0.13 nats more surprised than the chat model by the plain answer and no more surprised by the story, while the poet pays 0.56 nats on the plain answer. Verse pulls the words further from plain speech than pirate talk does; compare where the two curves above start.

There are several ways to limit it:

• A small rank and few steps. An adapter simply cannot move the weights far.
• A little of the old data in the new mix. A few percent of ordinary conversations remind the model of its usual manner; this is called replay.
• Stop in time. Watch the loss on held-out data of both kinds and stop when the old one starts to rise.
• Protect important weights. Elastic Weight Consolidation (Kirkpatrick et al., 2017) penalises changes to the weights that mattered most for the old task.

One more honest caveat. Fine-tuning, full or LoRA, is good at teaching manner: tone, format, length, a way of reasoning aloud. It is poor at adding knowledge, especially to a model as small as Sprout, whose 17 million numbers are already busy with English and simple stories. If you want Sprout to "know" your recipes, the realistic path is to show it the recipe right in the question.

Your own fine-tune

Everything above can be repeated with your own data. Here is the whole recipe.

1. Data. A file with one conversation per line, in the same format as in chapter 12:

{"messages": [{"role": "user", "content": "how do birds fly?"}, {"role": "assistant", "content": "Arr, they flap their wings and push the air down, matey, and the air pushes them up, just like wind in a sail!"}]}

A few hundred conversations are enough for a style; our adapters had 500 each. Consistency matters more than volume: if half the answers are pirate and half are not, the adapter learns to be half a pirate. Keep the answers at Sprout's level, short and in simple words. Five hundred examples won't teach it to talk like a professor: it barely knows the words.

2. Training. We start from the chat model, since a style sits on top of the ability to talk:

# train the adapter and merge it into the weights
python lora.py --base runs/chat/ckpt_final.pt --chats my-style.jsonl --out runs/my-style
# talk to the result in the terminal
python chat.py runs/my-style/ckpt_final.pt
# pack it for the browser: int8, as in chapter 11
python export.py runs/my-style/ckpt_final.pt sprout-my-style.bin

On a Mac with Apple silicon, lora.py uses the GPU through MPS automatically, and with an NVIDIA card it uses CUDA. In Google Colab a free GPU runtime is enough: upload the scripts, data/tokenizer.json and the chat checkpoint. Sprout is so small that the adapter trains even on a plain CPU, just more slowly. The curves widget above shows how long our own run took.

3. Checking. Watch the validation loss in log.jsonl: if it starts to rise while the training loss keeps falling, the adapter is memorising your examples instead of learning the style, so reduce --epochs. Talk to the result about topics that are not in your data: that's where forgetting shows. The knobs worth turning are --rank (compare 2, 8 and 32 by their validation loss), --lr and --epochs.

After merging, Pirate Sprout generates text in your browser. How fast is it compared with the plain chat model?

Slower: every token has to go through the detour B·A·x as well Exactly as fast: after merging it is the same model shape with different numbers Faster: LoRA removed most of the parameters

W′ = W + s·BA has the same shape as W. The engine runs the same matrix multiplications as before; only the numbers inside them are different.

Sprout right now

Sprout has learned to change its manner of speaking without relearning how to speak. Each style is a rank-8 correction of 417,792 numbers, trained on 500 conversations and merged back into the weights. Switch voices right in the middle of a conversation and see who remembers what. In the last chapter we'll gather every version of Sprout from this course in one garden and look at where to go next.