Fine-tuning and LoRA

You have a base model that writes fluent text, and you want it to do something more specific: answer in your support team’s style, follow instructions, or write SQL. Training a new model from scratch would cost millions and throw away everything the base model learned. Instead you keep training the model you have, on a much smaller set of examples. This lesson is about how to do that cheaply.

From pretraining to fine-tuning

The lesson on training described the stages. PretrainingPretrainingThe first and most expensive stage of training a language model, in which it learns to predict the next token across a very large amount of text. The result is a base model.Taught in How LLMs are trainedOpen in glossary on trillions of tokens produces a base model. Fine-tuningFine-tuningContinuing to train a pretrained model on a smaller, targeted dataset, such as example conversations, to change its behavior. Supervised fine-tuning uses examples of good responses.Taught in How LLMs are trainedOpen in glossary continues training on far less data chosen for a purpose, with the same next-token loss.

One especially important kind is instruction tuningInstruction tuningFine-tuning a pretrained model on many tasks written as instructions paired with good responses, so it learns to follow instructions it has not seen before.Open in glossary: fine-tuning on many tasks written as instructions paired with good responses. Wei et al. (2021) instruction-tuned a 137-billion-parameter model on more than 60 datasets phrased as instructions, and it got better at following instructions for kinds of tasks it had never seen. The supervised stage of InstructGPT (Ouyang et al., 2022) is the same idea with responses written by people.

The simplest approach, full fine-tuning, updates every weight. That is expensive in two ways:

• Memory while training. Training with Adam in mixed precision keeps, for each parameter, a 16-bit weight, a 16-bit gradient, and 12 bytes of optimizer state (a 32-bit copy of the weight plus two running averages), 16 bytes in all (Rajbhandari et al., 2020). For Llama 2 7B, with 6.74 billion parameters, that is about 108 GB before counting activations, more than a single 80 GB GPU holds.
• Storage afterwards. Every fine-tuned version is a complete copy of the model: 13.5 GB at 16 bits for each task or customer.

Learn the change, not the model

Write the fine-tuned weights of one layer as the pretrained weights plus a change: $W = W_0 + \Delta W$. Full fine-tuning learns $\Delta W$ with as many numbers as $W_0$ has. Low-rank adaptationLow-rank adaptationA way to fine-tune a model by freezing its weights and learning each change as the product of two thin matrices, B times A, of a small rank r. It trains a tiny fraction of the parameters.Open in glossary (LoRA, Hu et al., 2021) freezes $W_0$ and learns $\Delta W$ as the product of two thin matrices:

where $W_0$ is $d \times k$ and the rank $r$ is small, often between 1 and 64. Only $B$ and $A$ are trained, so one adapted matrix trains $r(d + k)$ numbers instead of $dk$. For a $4096 \times 4096$ matrix at $r = 8$, that is 65,536 instead of 16.8 million.

Three details matter:

• $B$ starts at zero. $A$ starts with small random numbers and $B$ with zeros, so $BA = 0$ and training begins exactly at the pretrained model. ($A$ must be random: if both started at zero, neither would ever get a gradient.)
• The scale $\alpha / r$. The update is multiplied by a constant $\alpha$ divided by the rank. Hu et al. set $\alpha$ once and did not tune it; dividing by $r$ keeps the size of the update from changing much when you try a different rank, so other settings such as the learning rate need less retuning.
• No extra cost when serving. After training, $\frac{\alpha}{r}BA$ can be added into $W_0$, giving an ordinary weight matrix. The model runs exactly as fast as before. Or you can keep one base model and swap small adapters per task.

The product $BA$ can never have rankRank (of a matrix)The number of independent directions a matrix can produce. A matrix of rank r can be written as a sum of r pieces, each one column times one row.Open in glossary higher than $r$: every change it can express is built from at most $r$ independent directions. Hu et al. hypothesized that the changes fine-tuning makes have a low “intrinsic rank”, building on earlier evidence that fine-tuning works even when restricted to a small random subspace of the parameters (Aghajanyan, Zettlemoyer, and Gupta, 2020).

How much can rank r capture?

The singular value decompositionSingular value decompositionWriting a matrix as a sum of rank-one pieces, each with a weight called a singular value, largest first. Keeping the largest r pieces gives the best rank-r approximation.Open in glossary answers that exactly. It writes any matrix as a sum of rank-one pieces, each with a weight called a singular value, largest first. The best rank-$r$ approximation keeps the $r$ largest pieces and drops the rest (the Eckart-Young theorem), so the singular values tell you, before any training, how much of a change each rank can capture.

The demo uses a $32 \times 32$ target change. The middle panel shows the best $BA$ for your rank, the right panel what it misses. Train B and A runs gradient descent from $B = 0$, as LoRA does, on the squared difference from the target.

import numpy as np

rng = np.random.default_rng(0)
d = 32

# A target change built like the demo's: four strong directions plus faint noise.
U, _ = np.linalg.qr(rng.normal(size=(d, d)))
V, _ = np.linalg.qr(rng.normal(size=(d, d)))
s = np.full(d, 0.15)
s[:4] = [8, 5, 3, 2]
delta = U @ np.diag(s) @ V.T

sv = np.linalg.svd(delta, compute_uv=False)
for r in (1, 2, 4, 8):
    kept = (sv[:r] ** 2).sum() / (sv ** 2).sum()  # the best rank-r share
    print(f"rank {r}: B and A hold {r * (d + d):4d} numbers (full: {d * d}), "
          f"best possible captures {kept:.1%}")

# LoRA starts exactly at the base model, because B starts at zero.
W = rng.normal(size=(d, d))
r, alpha = 4, 8
A = rng.normal(0, 1 / np.sqrt(d), size=(r, d))
B = np.zeros((d, r))
x = rng.normal(size=d)
print("starts at the base model:", np.allclose(W @ x + (alpha / r) * B @ (A @ x), W @ x))

# Real shapes: GPT-3 175B adapts W_q and W_v (12,288 x 12,288) in 96 layers.
gpt3 = 96 * 2 * 4 * (12288 + 12288)
print(f"GPT-3, r = 4 on W_q and W_v: {gpt3:,} numbers, {2 * gpt3 / 1e6:.1f} MB at 16 bits")

rank 1: B and A hold   64 numbers (full: 1024), best possible captures 62.4%
rank 2: B and A hold  128 numbers (full: 1024), best possible captures 86.7%
rank 4: B and A hold  256 numbers (full: 1024), best possible captures 99.4%
rank 8: B and A hold  512 numbers (full: 1024), best possible captures 99.5%
starts at the base model: True
GPT-3, r = 4 on W_q and W_v: 18,874,368 numbers, 37.7 MB at 16 bits

What you saw is the whole bet behind LoRA. When a change has a few dominant directions, as Low-rank update does, a small rank captures nearly all of it: rank 4 keeps 99.4% with a quarter of the numbers. When it has none, as Random noise does, rank 17 is needed for 90%, and at that rank $B$ and $A$ hold more numbers than the matrix itself.

The demo simplifies in one important way. It fits a known target change directly. Real LoRA never sees the change it is aiming for; it follows the gradient of the model’s loss on training examples, through every layer. The rank still limits what it can learn in exactly the way the singular values show.

Real numbers

Hu et al. applied LoRA to GPT-3 175B, adapting only the query and value projections, $W_q$ and $W_v$, in all 96 layers. Each is $12{,}288 \times 12{,}288$, so at $r = 4$ the adapter trains $96 \times 2 \times 4 \times (12{,}288 + 12{,}288) = 18{,}874{,}368$ numbers, about one ten-thousandth of the model’s 174.6 billion. They report that this cut the checkpoint from 350 GB to about 35 MB and training memory from 1.2 TB to 350 GB, and that LoRA matched or beat full fine-tuning on the benchmarks they tested.

The frozen weights still have to sit in memory, 13.5 GB at 16 bits for Llama 2 7B, but they need no gradients and no optimizer state. Storing the frozen weights in 4 bits shrinks that too; that combination, QLoRA, comes up in the lesson on quantization.