Your own fine-tune

We have a chat model that replies in simple sentences, tells stories and sometimes acknowledges its limitations. Now let's change its manner of speaking: add pirate expressions, for instance, or make it answer in rhyme. We would like to preserve useful abilities and avoid storing a whole new model for every style. How can we make this adjustment compact?

We could allow every weight to change again, as in the last chapter. Full fine-tuning, however, needs memory for gradients and optimizer state, and its result is usually saved as another set of weights. That is manageable for tiny Sprout; with billions of parameters, the costs become substantial. There is a way to train and store far fewer additional numbers.

That approach is LoRA, low-rank adaptation. It freezes the original weights during training and adds small trainable corrections to selected linear layers. Whether those corrections are sufficient is something we must test. First, we will use pictures to explore matrix rank and see why two small tables can sometimes stand in for a large one.

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.

To estimate memory, imagine full fine-tuning with Adam, storing weights, gradients and both running averages in float32. That is four four-byte numbers, or 16 bytes per parameter: about 277 MB for Sprout and 128 GB for an eight-billion-parameter model, before activations and temporary buffers. Actual usage depends on the optimizer and formats. Muon has different state, while mixed precision, quantisation and distributing work across devices change the calculation.

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?

Most of Sprout's parameters belong to matrices, rectangular tables of numbers; normalisation also has one-dimensional gains. For example, the combined query, key and value projection maps 384 input numbers to 1,152 outputs and contains 442,368 weights. To get a feel for matrix structure, imagine a picture where each entry controls the brightness of one pixel.

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. If both vectors are nonzero, the 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.

Fine-tuning changes a weight matrix by $\Delta W = W_\text{after} - W_\text{before}$. Could a useful correction have low rank? Studies of the intrinsic dimension of fine-tuning found that some tasks can be learned by moving within a much smaller parameter space. Low intrinsic dimension and low rank of an individual matrix are different ideas, though. LoRA makes a practical choice: restrict the search to low-rank corrections and test whether they suffice for the new task.

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$ sets the scale of the adapter's contribution. Our lora.py uses $\alpha=16$ and $r=8$, multiplying the correction by 2. Dividing by rank helps control scale but does not guarantee equivalent behaviour for every $r$. After changing rank, we should still check learning rate and quality. We will also mention rsLoRA below, which uses a different scaling rule.

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$.

The pictures show why different changes need different ranks. The poet's beret and scarf consist mainly of horizontal and vertical stripes. The pirate's diagonal eye-patch strap is harder to express as a few outer products: its dark point moves across columns from row to row. At rank 8, the pirate's error is noticeably higher. The dashed line gives the best attainable error at that rank, from the Eckart–Young theorem. Compare it with the optimizer's result; how closely it approaches the bound within the allotted steps depends on the settings.

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 mostly convenient: small adapters are easy to store and swap. With a large model, the memory savings matter more. Eight billion frozen weights in a 16-bit format occupy roughly 16 GB; adapters, their gradients and optimizer state add comparatively little. Activations and working buffers still need memory, sometimes a great deal. QLoRA reduces the storage for base weights through 4-bit quantisation. Whether the complete run fits on a laptop also depends on context length, batch size, rank and implementation.

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.

One way to detect a loss of an old ability is to measure prediction error again on texts of the previous kind. Give four versions of Sprout the same examples and compare average $-\ln p$ per token. A higher loss means a version predicts that particular text less well. A handful of short examples offers a useful clue, but assessing an entire skill requires a broader evaluation.

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 measures worse prediction of this example, not all previous abilities at once. 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. This limits correction complexity and training duration, not magnitude: even a rank-one matrix can contain very large numbers.
• 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.

A family recipe poses a different problem from a pirate voice: preserving the exact ingredients and amounts matters. Fine-tuning can teach new information, but does not guarantee faithful recall, particularly with a small model and a few examples. For an experiment, it is easier to include a short recipe in the question and check whether Sprout uses it correctly. That still does not promise a flawless answer, but it makes the source available to both model and reader.

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!"}]}

Each adapter in our experiment used 500 conversations. That is a starting point for a similar style, not a minimum for every task. Varied questions and consistent answers matter: adding “arr” to every response does not itself teach the model to address the question. Start with short sentences and vocabulary familiar to Sprout, leaving some topics for evaluation. If you mix the new style with ordinary speech, check when and how the model uses each.

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

lora.py selects an available device: CUDA for supported NVIDIA hardware, MPS for Apple silicon, or the CPU. You need the scripts, data/tokenizer.json, the chat checkpoint and your conversation file. Sprout is small enough to experiment with without a large server, but runtime depends on the device and settings. The curves above report the duration of our two runs.

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 The same operation count per token: merging changes the weights, not the model shape 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.