Fine-tuning and LoRA

Problem

Pretraining spends enormous compute to produce one general model. Every use of it afterwards, a support bot, a code reviewer, a classifier for legal clauses, needs a different behaviour. Training a new model per task from scratch throws the pretraining away. Continuing to train the whole model per task works, but each copy of a 7-billion-parameter decoder is about 13.5 GB in 16-bit, and training it with Adam needs several times that in GPU memory.

Two further problems come with training everything. A small task dataset can pull the weights far from where pretraining left them, and the model forgets what it knew: catastrophic forgetting. And a model that was only taught to continue text does not follow instructions; that has to be taught too. The question is how to adapt a large model cheaply, without breaking it, and keep many adaptations side by side.

Intuition

Think of the pretrained weights as an experienced employee. Feature extraction leaves them untouched and trains only a small head on their outputs, like asking for a report in a new format. Full fine-tuning retrains the whole person, which is powerful but expensive and can overwrite old skills. The parameter-efficient methods sit between: they freeze the original and train a small add-on.

LoRA's bet is that the change needed for a new task is simple even when the weights are not. A weight matrix in a 4096-wide model has 16.8 million entries, but the useful update may move it along only a few directions. So LoRA writes the update as a product of two thin matrices, B times A, with an inner dimension (the rank) of 8 or so. That is the low-rank approximation idea from the SVD, used not to compress a matrix but to restrict what training can change. Since B starts at zero, the model starts exactly where pretraining left it.

Mechanics

Full fine-tuning versus feature extraction

Feature extraction (a linear probe when the head is linear) freezes the network and trains only a new output layer on its final hidden states. It is cheap, cannot forget, and needs no backward pass through the body, but it can only use features that already exist. Full fine-tuning updates every weight, usually with a learning rate well below the pretraining one so the weights stay near their starting point; it is the most accurate when data is plentiful and the most expensive. In between, one can unfreeze only the top layers. Kumar et al. found that full fine-tuning with a randomly initialized head can distort the pretrained features while the head is still bad, and that training the linear probe first and then fine-tuning everything (LP-FT) holds up better out of distribution.

In PyTorch, freezing is module.requires_grad_(False), which sets requires_grad on every parameter; frozen tensors get no gradient and should be left out of the optimizer, so they cost no optimizer state.

Instruction tuning

A pretrained language model continues text; asked a question, it might continue with three more questions. Instruction tuning, also called supervised fine-tuning (SFT), trains it on pairs of an instruction and a good response, formatted with the chat template the model will see at inference. Wei et al. (FLAN) showed that fine-tuning on many tasks phrased as instructions improves zero-shot performance on tasks never seen in training, and InstructGPT used SFT on human-written demonstrations as the first stage before preference optimization.

The loss is the usual next-token cross-entropy, but only on the response tokens: prompt positions get the label −100, which is F.cross_entropy's default ignore_index, so they add nothing to the loss or its average. The model is trained to answer, not to reproduce the user's question.

Catastrophic forgetting

Gradient descent on task B moves the weights toward B's optimum with no regard for task A, and nothing in B's loss says A mattered. Remedies all keep the weights near the pretrained point: a small learning rate and few epochs; an L2 penalty toward the original weights; elastic weight consolidation (EWC, Kirkpatrick et al.), which weights that penalty by each parameter's Fisher information so important weights move least; mixing some original-style data into the fine-tuning set; and freezing most weights, which is what the parameter-efficient methods do structurally. Biderman et al. found that LoRA learns less of a new domain than full fine-tuning but also forgets less of the old one.

LoRA, derived

Take one linear layer with frozen weight W₀ ∈ ℝ^d×k (PyTorch's nn.Linear stores it as [out_features, in_features] = [d, k]). Full fine-tuning learns an update ΔW of the same shape: dk numbers. LoRA constrains it to rank at most r by writing it as a product:

ΔW = (α / r) · B A,   B ∈ ℝ^d×r,   A ∈ ℝ^r×k,   r ≪ min(d, k)
h = W₀x + (α / r) · B(A x)

Every column of BA is a combination of B's r columns, so its rank is at most r. The trainable count drops from dk to r(d + k). For a 4096 × 4096 attention projection at r = 8: 16,777,216 becomes 8 · 8192 = 65,536, 0.39% as many.

A starts random (Gaussian in the paper) and B starts at zero, so ΔW = 0 and training begins at the pretrained model; both cannot be zero, or no gradient would flow to either. The constant α / r scales the update so that changing r does not require retuning the learning rate much. Computing B(Ax) right to left never forms the d × k product. After training, W = W₀ + (α/r)BA can be merged into a single matrix, so inference has no extra latency, and a different adapter can be swapped in by subtracting one product and adding another.

Hu et al. applied LoRA to the attention query and value projections and found very small ranks competitive; on GPT-3 175B it cut trainable parameters about 10,000-fold and training GPU memory about threefold relative to full fine-tuning with Adam. Later practice, including QLoRA, often applies it to every linear layer.

QLoRA

With LoRA, the frozen base still has to sit in GPU memory, 2 bytes per parameter in 16-bit. QLoRA (Dettmers et al.) stores the frozen base in 4 bits and trains LoRA adapters on top, backpropagating through the quantized weights, which are dequantized to bfloat16 for each matmul. Three pieces make it work:

• 4-bit NormalFloat (NF4): a 4-bit data type whose 16 levels are quantiles of a normal distribution, matched to how pretrained weights are distributed. Each block of 64 weights shares one absmax scale (quantization).
• Double quantization: the per-block scales are themselves quantized to 8 bits in blocks of 256, cutting their overhead from 32/64 = 0.5 bits per parameter to about 0.127.
• Paged optimizers: optimizer state in unified memory that pages to CPU RAM during memory spikes.

The paper fine-tuned a 65B model on a single 48 GB GPU, and found adapters on all linear layers were needed to match 16-bit full fine-tuning.

Adapters, prefix tuning and prompt tuning

Adapters (Houlsby et al.) insert a small bottleneck MLP, down-project d → m, nonlinearity, up-project m → d, with a residual connection, twice in every transformer layer, and train only those. They are near-identity at initialization like LoRA, but they add layers to the forward pass, so they add inference latency and cannot be merged away.

Prefix tuning (Li and Liang) prepends a few trainable vectors to the keys and values of every attention layer: virtual tokens that every real token can attend to. Prompt tuning (Lester et al.) keeps only trainable embeddings at the input, a soft prompt; it becomes competitive with full fine-tuning as models grow to around ten billion parameters. Both use up context positions and are often less stable to train than LoRA.

Worked example

All numbers below come from tested helpers in ml_models_docs.py (ft_lora_toy, ft_model_params, ft_lora_trainable, ft_memory, ft_masked_loss, ft_forgetting). The Llama 2 7B count matches its published size, and the LoRA counts match what the Hugging Face PEFT library reports for the same settings.

Implementation

LoRA with torch.nn.utils.parametrize: the parametrization receives the frozen weight and returns W₀ + (α/r)BA, so nn.Linear needs no changes, and remove_parametrizations performs the merge.

import torch
import torch.nn as nn
import torch.nn.utils.parametrize as P


class LoRA(nn.Module):
    def __init__(self, out_features, in_features, r=8, alpha=16):
        super().__init__()
        self.A = nn.Parameter(torch.randn(r, in_features) * 0.01)
        self.B = nn.Parameter(torch.zeros(out_features, r))
        self.scale = alpha / r

    def forward(self, W):
        return W + self.scale * (self.B @ self.A)


def add_lora(model, names=("q_proj", "v_proj"), r=8, alpha=16):
    model.requires_grad_(False)
    targets = [m for n, m in model.named_modules() if isinstance(m, nn.Linear) and n.endswith(names)]
    for layer in targets:
        P.register_parametrization(layer, "weight", LoRA(layer.out_features, layer.in_features, r, alpha))
    return [p for p in model.parameters() if p.requires_grad]


def merge_lora(model):
    for m in model.modules():
        if P.is_parametrized(m, "weight"):
            P.remove_parametrizations(m, "weight", leave_parametrized=True)


trainable = add_lora(model)
optimizer = torch.optim.AdamW(trainable, lr=2e-4)
# ... train on (input_ids, labels) with prompt positions set to -100 ...
merge_lora(model)

The original weight moves to layer.parametrizations.weight.original and keeps requires_grad=False, so only A and B reach the optimizer. This version forms the full d × k weight on every forward pass (wrap the forward in P.cached() to do it once per pass when a weight is used more than once); production libraries such as PEFT instead wrap the layer and compute base(x) + scale * (x @ A.T @ B.T), which never builds the large product. The SFT loss is F.cross_entropy(logits.view(-1, V), labels.view(-1)), whose default ignore_index=-100 drops the prompt.

Tradeoffs

• Full fine-tuning has the most capacity and is the reference for quality on large, shifted datasets, at full memory cost, one full model copy per task, and the most forgetting.
• Feature extraction is cheapest and cannot forget, but cannot learn features the base lacks.
• LoRA trains a fraction of a percent, stores a few megabytes per task, merges to zero inference cost, and forgets less; it can trail full fine-tuning when the new task is far from pretraining. Rank, α and which layers to target are the knobs.
• QLoRA cuts base-model memory about fourfold again, at the cost of dequantization work every step and an adapter trained against quantized weights; merging back into a 16-bit model is then an approximation.
• Adapters are modular but add latency; prefix and prompt tuning are tiny but use context positions and train less reliably, especially on smaller models.
• Instruction data quality matters more than quantity: SFT copies the style and the mistakes of its demonstrations.

Connections

References

• Hu et al. LoRA: Low-Rank Adaptation of Large Language Models. ICLR 2022. arXiv:2106.09685
• Dettmers et al. QLoRA: Efficient Finetuning of Quantized LLMs. NeurIPS 2023. arXiv:2305.14314
• Houlsby et al. Parameter-Efficient Transfer Learning for NLP. ICML 2019. arXiv:1902.00751
• Li and Liang. Prefix-Tuning: Optimizing Continuous Prompts for Generation. ACL 2021. arXiv:2101.00190
• Lester, Al-Rfou and Constant. The Power of Scale for Parameter-Efficient Prompt Tuning. EMNLP 2021. arXiv:2104.08691
• Wei et al. Finetuned Language Models Are Zero-Shot Learners (FLAN). ICLR 2022. arXiv:2109.01652
• Ouyang et al. Training language models to follow instructions with human feedback (InstructGPT). NeurIPS 2022. arXiv:2203.02155
• Kirkpatrick et al. Overcoming catastrophic forgetting in neural networks. PNAS 2017. arXiv:1612.00796
• Kumar et al. Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distribution. ICLR 2022. arXiv:2202.10054
• Biderman et al. LoRA Learns Less and Forgets Less. TMLR 2024. arXiv:2405.09673
• Touvron et al. Llama 2: Open Foundation and Fine-Tuned Chat Models. 2023. arXiv:2307.09288
• Rajbhandari et al. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. SC 2020. arXiv:1910.02054
• PyTorch v2.15.0: Linear, register_parametrization, remove_parametrizations, cached, cross_entropy.