How LoRA works, written by hand

Last lesson said that when fine-tuning really is needed, LoRA is the first choice. This lesson works out what it is, and implements it by hand on the small GPT trained in Module 09.

python lora_from_scratch.py

The trouble with full fine-tuning

Fine-tuning means continuing to train an already trained model on new data. The most direct way is to train all the parameters, which is called full fine-tuning.

For an LLM that's expensive. During training, each parameter needs its gradient stored as well as itself, and the AdamW optimiser stores two extra numbers per parameter (Module 08, Lesson 5 mentioned that it records each parameter's "momentum"). For a 7-billion-parameter model, that alone takes over a hundred GB of GPU memory.

And every fine-tune produces a complete new model. Fine-tune a version for each of ten customers and you're storing ten files of tens of GB each.

The idea of LoRA

The idea of LoRA (Low-Rank Adaptation) is that the change a weight needs during fine-tuning doesn't actually need to be that "complex".

A linear layer's weight is a matrix W, say 128×384. Full fine-tuning learns a change ΔW of the same size and adds it to W. LoRA doesn't learn ΔW directly; it splits it into the product of two very thin matrices:

      ΔW       =       A       ×     B
  (128 × 384)      (128 × 8)     (8 × 384)

  49152 个数       1024 个数   +  3072 个数  =  4096 个数

The 8 in the middle is called the rank (r). A and B together have only 4,096 numbers, a twelfth of ΔW. The smaller r is, the fewer parameters to learn, but the more limited the changes it can express.

During training, the original W is frozen, and only A and B are trained:

输出 = x × W  +  x × A × B × 缩放系数
       ─────     ─────────────────────
      原来的路        新加的支路

Implementation

Wrapping a trained nn.Linear takes only a dozen or so lines:

class LoRALinear(nn.Module):
    """包住一个已经训练好的线性层：原来的权重冻结不动，旁边加一条 A、B 两个小矩阵的"支路"。"""

    def __init__(self, base: nn.Linear, rank=8, alpha=16):
        super().__init__()
        self.base = base
        for p in self.base.parameters():
            p.requires_grad = False  # 原来的参数一个都不训练
        self.A = nn.Parameter(torch.randn(base.in_features, rank) * 0.01)  # 输入维度 × r
        self.B = nn.Parameter(torch.zeros(rank, base.out_features))  # r × 输出维度，初始为 0
        self.scale = alpha / rank

    def forward(self, x):
        # 原来的输出 + 支路的输出。B 一开始是 0，所以刚加上时模型的行为和原来完全一样
        return self.base(x) + (x @ self.A @ self.B) * self.scale

Two details.

B is initialised to 0. That way A×B is 0 at the start, the side branch has no effect, and the model is exactly the same as before. Training starts from the original model, not from one scrambled by a random side branch. A mustn't also be 0, or the gradients of both A and B would be 0 and nothing would ever be learned.

The scaling factor alpha / r. It controls how much influence the side branch has. With it, you don't need to retune the learning rate when you change r. r=8 with alpha=16 is a very common setting.

Then replace every attention layer's qkv and proj in the model with LoRA:

for p in model.parameters():
    p.requires_grad = False
for block in model.blocks:
    block.attn.qkv = LoRALinear(block.attn.qkv)
    block.attn.proj = LoRALinear(block.attn.proj)

The task: make it write only five-character quatrains

The small GPT trained in Module 09 writes poems in every format. Now fine-tune it on five-character quatrains alone and see whether it can be made to write only five-character quatrains.

原来的训练数据里，五言绝句占 10.4%（3538 首）
微调前：生成 200 首，其中五言绝句 67 首

模型共 1,639,296 个参数，要训练的只有 24,576 个（1.50%）
刚装上 LoRA 时，输出和原模型一样吗？True

Before fine-tuning, 67 of 200 poems are five-character quatrains. That's quite a bit higher than the 10.4% in the training data, perhaps because five-character quatrains are the shortest and easiest to write completely.

With LoRA added, only 24,576 parameters are trained, 1.5% of the whole model. 4 layers, two LoRAs per layer: qkv's A and B are 128×8 and 8×384, proj's are 128×8 and 8×128, 6,144 per layer, 24,576 for 4 layers.

"Is the output the same right after adding LoRA? True" confirms the effect of initialising B to 0.

Training

Using only five-character quatrains as data, train for 300 steps:

  第 1 步  损失 4.356
  第 100 步  损失 4.294
  第 200 步  损失 4.335
  第 300 步  损失 4.213
训练 300 步用了 23 秒

微调后：生成 200 首，其中五言绝句 191 首。前 5 首：
  白头还作雨，一鬓自如霜。何时之饮在，试死势悠扬。
  白发勤鸿急，秋贫梦不闲。夜寒休绕郡，秋杀杜陵陂。
  才子俭高人，何言亦重重。言知不见人，辄空一一言。
  一丛生碧霄，危棹度清湍。沙渚穿花发，苍苍隔鹤闲。
  九陌两金阁，一枝归草间。何当重枕簟，归复独裴回。

191 of 200 poems are five-character quatrains, up from 33.5% to 95.5%. Only 1.5% of the parameters were touched, and training took 23 seconds.

Interestingly, the loss barely changed, from 4.36 to 4.21. This shows the model's poetry "skill" (how accurately it predicts each character) didn't improve; what changed was its "habit": after writing four five-character lines, it learned to output a newline and end the poem. That's exactly the kind of change LoRA is good at: adjusting format, style and behaviour, not teaching the model lots of new knowledge. Last lesson said "fine-tuning changes behaviour; knowledge comes from RAG", and here you can see it clearly.

Saving and merging

LoRA 的权重单独存下来只有 96 KB，整个模型是 6.3 MB
把 LoRA 合并回原来的权重之后，模型又变回 1,614,720 个参数，输出和合并前最大差别 5.7e-06

LoRA's first benefit: you only need to store A and B. Here that's 96 KB, while the model itself is 6.3 MB. With LLMs the gap is even starker: for a model of tens of GB, the LoRA weights may be only tens of MB. Fine-tuning once for each of ten customers means one base model plus ten small files.

The second benefit: it can be merged. After training, add A×B×scaling factor directly to the original W, and you get an ordinary linear layer:

def merged(self):
    """把支路合并回原来的权重，得到一个普通的线性层：推理时没有任何额外开销。"""
    layer = nn.Linear(self.base.in_features, self.base.out_features)
    with torch.no_grad():
        # nn.Linear 的权重形状是 (输出, 输入)，所以要转置
        layer.weight.copy_(self.base.weight + (self.A @ self.B).T * self.scale)
        layer.bias.copy_(self.base.bias)
    return layer

After merging, the model's structure and parameter count are exactly the same as before, and inference isn't any slower. The output differs from before merging by 5.7e-06, just floating-point error.

Why LoRA works

An intuitive explanation: a pre-trained model has already learned most things, and fine-tuning just makes some "small adjustments" on top. The 2021 paper that introduced LoRA (Hu et al., "LoRA: Low-Rank Adaptation of Large Language Models") starts from exactly this point: the change needed for such adjustments is itself "low-rank", and can be expressed with a very small r.

Our experiment fits this explanation too: making the model "stop after four lines" is a very simple change in behaviour, and 1.5% of the parameters is plenty. But to teach a model lots of entirely new knowledge, such as the content of a whole new field, LoRA won't necessarily match full fine-tuning.

LoRA has a commonly used variant called QLoRA: quantise the frozen original model to 4 bits (Lesson 4 covers quantisation) to save GPU memory, while LoRA's A and B are still trained at normal precision. This makes it possible to fine-tune models with billions of parameters on a single consumer graphics card.

Exercises

1. Change rank to 1, 2 and 32. What share of poems are five-character quatrains after fine-tuning in each case, and how many parameters does each have?
2. Skip LoRA and do full fine-tuning (train all parameters, with the learning rate changed to 1e-4) for the same 300 steps, and compare the share of five-character quatrains and the training time.
3. Add LoRA only to the feed-forward networks (block.mlp[0] and block.mlp[2]), not to attention. How do the results differ?

Self-check

1. Which parameters does LoRA train, and what happens to the original parameters?

All the original weights are frozen and don't take part in training. LoRA adds two small matrices, A and B, alongside the chosen linear layers and trains only them. The output equals the original output plus x×A×B×scaling factor.

2. Why is B initialised to 0? Could A also be initialised to 0?

With B at 0, A×B is 0, so a model with LoRA just added is exactly the same as the original, and training starts from the original model. A can't also be 0: if both A and B were 0, both gradients would be 0 too, and the parameters would never update.

3. After LoRA training, is inference slower?

It needn't be. After training, add A×B×scaling factor to the original weights to merge them, and the resulting model has the same structure and parameter count as the original, and the same inference speed. Without merging, there's an extra side branch to compute, which is a little slower, but the benefit is that you can switch between different LoRAs at any time.