Module 3 · Reading 3.6 · read before S19: Fine-Tuning & LoRA

The PDF posted in Canvas is the assigned version of this reading. It is built from this page, so the text is identical; anything here marked as deep reading is optional and is not in the PDF.

Pretraining gives you a model that knows a great deal about language in general and nothing in particular about your task. The classical remedy, familiar from transfer learning with CNNs, is fine-tuning: continue gradient descent on your task’s data. For a 10-million-parameter CNN this is routine. For a 7-billion-parameter language model it collides with arithmetic — not the cost of storing the model, but the cost of training it, which for Adam-family optimizers is roughly four times larger than the weights themselves. The result is that full fine-tuning of large models is out of reach on any single ordinary GPU, and out of budget for most teams even with clusters.

This section does that arithmetic explicitly, then develops the idea that broke the impasse: LoRA (Low-Rank Adaptation, Hu et al., 2021), which freezes the pretrained weights and trains only a pair of small matrices per layer. You will implement a LoRA layer from scratch, count its parameters, and verify numerically the property that makes it free at inference time: the learned update can be merged back into the frozen weight matrix.

How to read this

The cost of full fine-tuning

Fine-tuning runs the same training loop as pretraining, so it pays the same per-parameter memory bill. With standard fp32 training and Adam or AdamW, every trainable parameter drags along:

  • the weight itself: 4 bytes;
  • its gradient: 4 bytes;
  • Adam’s first moment m (the running mean of gradients): 4 bytes;
  • Adam’s second moment v (the running mean of squared gradients): 4 bytes.

That is 16 bytes per trainable parameter for the particular fp32/mixed-precision Adam layouts assumed here. Real optimizers differ: some omit a master copy, keep gradients in fp32, use 8-bit moments, shard states, or allocate states lazily. Treat 16 bytes as a planning recipe whose assumptions must be stated, not a universal constant.

# Memory arithmetic for full fine-tuning with Adam/AdamW.
# Per trainable parameter, a standard fp32 setup stores:
#   weight (4 bytes) + gradient (4) + Adam m (4) + Adam v (4) = 16 bytes.
# Common mixed-precision setup (bf16 weights/grads, fp32 master + states):
#   bf16 weight (2) + bf16 grad (2) + fp32 master copy (4) + m (4) + v (4) = 16 bytes.
# Activations come on top of this and depend on batch/sequence length.

GiB = 1024 ** 3

def full_ft_memory(n_params, bytes_per_param=16):
    return n_params * bytes_per_param / GiB

def lora_memory(n_params, n_trainable, weight_bytes=2):
    # frozen base weights (no grads, no optimizer states) + full stack for LoRA params
    return (n_params * weight_bytes + n_trainable * 16) / GiB

models = [("124M (GPT-2 small)", 124e6),
          ("1.3B", 1.3e9),
          ("7B", 7e9),
          ("70B", 70e9)]

print(f"{'model':>20} {'full FT (GiB)':>14} {'LoRA r=8 ~0.1% (GiB)':>21}")
for name, n in models:
    n_lora = 0.001 * n  # ~0.1% trainable, a typical LoRA budget
    print(f"{name:>20} {full_ft_memory(n):>14.1f} {lora_memory(n, n_lora):>21.1f}")

n = 7e9
print(f"\nbreakdown for the 7B model, full fine-tuning (fp32):")
print(f"  weights            : {n * 4 / GiB:6.1f} GiB")
print(f"  gradients          : {n * 4 / GiB:6.1f} GiB")
print(f"  Adam m (momentum)  : {n * 4 / GiB:6.1f} GiB")
print(f"  Adam v (variance)  : {n * 4 / GiB:6.1f} GiB")
print(f"  total (no activations): {full_ft_memory(n):.1f} GiB")

Running this prints:

               model  full FT (GiB)  LoRA r=8 ~0.1% (GiB)
  124M (GPT-2 small)            1.8                   0.2
                1.3B           19.4                   2.4
                  7B          104.3                  13.1
                 70B         1043.1                 131.4

breakdown for the 7B model, full fine-tuning (fp32):
  weights            :   26.1 GiB
  gradients          :   26.1 GiB
  Adam m (momentum)  :   26.1 GiB
  Adam v (variance)  :   26.1 GiB
  total (no activations): 104.3 GiB

Read the 7B row carefully: the weights alone are 26 GiB, which fits on a high-end GPU, but training state pushes the total past 104 GiB before activations. This is why “I can run inference on this model” does not imply “I can fine-tune it.” The right-hand column previews the punchline of this section: freeze the base model (2 bytes per parameter, no gradients, no optimizer states) and train roughly 0.1% of the parameters, and the same 7B model needs about 13 GiB — single-GPU territory. There is a second, quieter cost of full fine-tuning: every task produces a complete copy of the model. Ten fine-tuned variants of a 7B model are 260 GiB of nearly identical checkpoints.

Two clarifications before moving on. First, LoRA usually does not substantially reduce the saved activation term through the frozen backbone, so long sequences can still dominate memory. Some PEFT strategies that train only a head or avoid backpropagating through large portions of the network can reduce activation storage, so “no PEFT method” would be too broad. Second, sharding and lower-precision optimizer states are other ways to reduce the per-device training bill.

The low-rank idea

LoRA starts from a hypothesis inspired by Aghajanyan et al. (2021), who showed fine-tuning works inside small random subspaces: the weight changes that fine-tuning produces seem to have low effective rank — adaptation happens in a much smaller subspace than the full parameter space. If the useful part of the update \(\Delta W\) for a weight matrix \(W\) (shape \(d_{\text{out}} \times d_{\text{in}}\)) is approximately low-rank, why parameterize the whole matrix? Write the update as a product of two thin matrices:

\[W' = W + \frac{\alpha}{r}\, B A\]

where \(A\) has shape \(r \times d_{\text{in}}\), \(B\) has shape \(d_{\text{out}} \times r\), and the rank \(r\) is small: 1 to 64 in the paper’s experiments. \(W\) stays frozen; only \(A\) and \(B\) receive gradients. The parameter count for the update drops from \(d_{\text{out}} \times d_{\text{in}}\) to \(r\,(d_{\text{in}} + d_{\text{out}})\). For a \(4096 \times 4096\) attention projection with \(r = 8\), that is 65,536 trainable parameters instead of 16.8 million (0.4% of the original).

If “rank” feels abstract, anchor it in the SVD picture from linear algebra: any matrix can be written as a weighted sum of rank-1 outer products, and its rank is how many of those terms it genuinely needs. Saying the fine-tuning update is “approximately rank 8” says the change to a 16-million-entry weight matrix is well described by 8 directions of input being remapped to 8 directions of output — everything else the layer does passes through untouched. That is a strong structural bet about adaptation: the pretrained matrix stays arbitrarily rich; only the correction is constrained to be simple.

Three design details matter in practice. First, initialization: \(A\) starts as small random Gaussian noise and \(B\) starts as zeros, so \(BA = 0\) and the model’s behavior at step 0 is exactly the pretrained model — training starts from a known-good point rather than a perturbed one. Second, the scale factor \(\alpha/r\) decouples the update’s magnitude from the chosen rank, so you can change \(r\) without retuning the learning rate (the paper’s stated motivation). Third, placement: Hu et al. applied LoRA to the attention projection matrices (\(W_q\) and \(W_v\) in their main experiments) and reported that this was sufficient; later practice often adapts every linear layer.

The payoff the paper reports for GPT-3 175B: roughly a 10,000x reduction in trainable parameters versus full fine-tuning and about a 3x reduction in GPU memory during training, with fine-tuning quality on their benchmarks on par with, sometimes better than, full fine-tuning. Those are the paper’s numbers on their setup, not a universal law; on tasks far from the pretraining distribution, low-rank updates can fall short of full fine-tuning, and rank becomes a hyperparameter you must actually tune.

Historical note. Parameter-efficient adaptation predates LoRA. Houlsby et al. (2019) inserted small trainable “adapter” bottleneck layers between frozen Transformer sublayers, and prompt-based methods, prefix tuning (Li and Liang, 2021) and prompt tuning (Lester et al., 2021), kept the entire network frozen and optimized a handful of continuous “virtual token” embeddings instead. LoRA’s 2021 contribution was showing you could get adapter-quality results with no extra inference cost: adapters add layers that must run at serving time, while a merged LoRA update disappears into the weights. That single property (plus a simple implementation) made it the default within about two years, and the “PEFT” umbrella term (parameter-efficient fine-tuning) entered common use with Hugging Face’s library of the same name.

LoRA from scratch

The implementation is short enough to write and verify completely. The class below wraps an existing nn.Linear, freezes it, and adds the low-rank path; after training we merge \((\alpha/r)\,BA\) into the base weight and check that the merged model computes the same function.

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)

class LoRALinear(nn.Module):
    """A frozen nn.Linear plus a trainable low-rank update (alpha/r) * B @ A."""
    def __init__(self, base: nn.Linear, r=4, alpha=8):
        super().__init__()
        self.base = base
        for p in self.base.parameters():
            p.requires_grad = False          # freeze W and bias
        self.r, self.alpha = r, alpha
        self.A = nn.Parameter(torch.randn(r, base.in_features) * 0.01)
        self.B = nn.Parameter(torch.zeros(base.out_features, r))

    def forward(self, x):
        return self.base(x) + (self.alpha / self.r) * (x @ self.A.T @ self.B.T)

    def merged_weight(self):
        return self.base.weight + (self.alpha / self.r) * (self.B @ self.A)

# --- A small "pretrained" model, then adapt it to a new task with LoRA ------
d_in, d_hidden, d_out = 64, 128, 10
model = nn.Sequential(nn.Linear(d_in, d_hidden), nn.ReLU(),
                      nn.Linear(d_hidden, d_out))
# pretend this model is pretrained; wrap both Linears with LoRA
model[0] = LoRALinear(model[0], r=4, alpha=8)
model[2] = LoRALinear(model[2], r=4, alpha=8)

total = sum(p.numel() for p in model.parameters())
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"total params    : {total:,}")
print(f"trainable (LoRA): {trainable:,} ({100 * trainable / total:.2f}%)")

# synthetic classification task
g = torch.Generator().manual_seed(1)
X = torch.randn(512, d_in, generator=g)
w_true = torch.randn(d_in, d_out, generator=g)
y = (X @ w_true).argmax(dim=1)

opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=1e-2)
for step in range(201):
    loss = F.cross_entropy(model(X), y)
    opt.zero_grad(); loss.backward(); opt.step()
    if step % 100 == 0:
        acc = (model(X).argmax(1) == y).float().mean()
        print(f"step {step:3d}  loss {loss.item():.4f}  acc {acc.item():.4f}")

# --- Merge check: replacing base weight with W + (alpha/r) B @ A ------------
with torch.no_grad():
    merged = nn.Sequential(nn.Linear(d_in, d_hidden), nn.ReLU(),
                           nn.Linear(d_hidden, d_out))
    for i in (0, 2):
        merged[i].weight.copy_(model[i].merged_weight())
        merged[i].bias.copy_(model[i].base.bias)
    x_test = torch.randn(64, d_in, generator=g)
    diff = (model(x_test) - merged(x_test)).abs().max()
print(f"max |LoRA model - merged model| on fresh inputs: {diff.item():.2e}")

Running this prints:

total params    : 10,930
trainable (LoRA): 1,320 (12.08%)
step   0  loss 2.3034  acc 0.0977
step 100  loss 0.0492  acc 1.0000
step 200  loss 0.0058  acc 1.0000
max |LoRA model - merged model| on fresh inputs: 2.34e-05

Walk through the three results. The parameter count: with r = 4, the trainable fraction on this tiny model is 12%, much higher than the fractions quoted for large models, because LoRA’s savings ratio \(r\,(d_{\text{in}} + d_{\text{out}}) / (d_{\text{in}} \cdot d_{\text{out}})\) shrinks as layers get wider. At \(d_{\text{in}} = d_{\text{out}} = 4096\) the same r = 4 gives 0.2%; parameter-efficiency is a big-model phenomenon. The training run: the frozen base (random, standing in for a pretrained one) plus two rank-4 updates reaches 100% training accuracy on this synthetic task, from a starting loss of \(2.3034 \approx \ln(10)\), chance level for 10 classes. The merge check: after copying \(W + (\alpha/r)\,BA\) into a plain nn.Linear, the merged network matches the wrapped network to within 2.3e-05 on fresh inputs. That residue is fp32 round-off, not an implementation error — the two computations associate the matrix products differently (\((xA^\top)B^\top\) versus \(x(BA)^\top\)), which changes floating-point rounding but not the underlying linear map.

The merge property is worth stating as the section’s central engineering fact: after training, a LoRA adapter costs zero additional inference latency, because it stops existing as a separate module. And because the adapter itself is tiny (kilobytes to megabytes), you can store one base model plus a shelf of task adapters, and swap tasks by swapping adapters, which also means you can unmerge (subtract the update) to recover the base model exactly, up to round-off.

People. LoRA came out of Microsoft Research in 2021, from a team including Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen — a mix of systems-minded researchers and learning theorists, working on the practical problem of serving many fine-tuned variants of GPT-3. Tim Dettmers, then a PhD student at the University of Washington known for pushing quantization (8-bit optimizers, bitsandbytes), led the 2023 QLoRA work that combined a 4-bit frozen base model with LoRA adapters; the paper reports fine-tuning a 65B model on a single 48 GB GPU. Much of the reason these methods spread quickly is that both teams released usable code.

The PEFT family

LoRA is one member of a family of parameter-efficient fine-tuning (PEFT) methods, and knowing the neighbors helps you read the literature:

  • Adapters (Houlsby et al., 2019): small bottleneck MLPs inserted after attention and feed-forward sublayers; the frozen network runs through them at inference, adding latency. The bottleneck design (project down to a few dimensions, apply a nonlinearity, project back up, add residually) is the same shape-shrinking trick as LoRA, but placed in series with the network rather than in parallel with a weight matrix, which is exactly why it cannot be merged away.
  • Prefix tuning (Li and Liang, 2021) and prompt tuning (Lester et al., 2021): keep every weight frozen and learn continuous embeddings prepended to the input (prompt tuning) or to each layer’s keys and values (prefix tuning). These “virtual tokens” are not words: they are free vectors in embedding space, optimized by gradient descent, and they steer the frozen computation the way a carefully chosen text prompt would, but without being constrained to the vocabulary. Lester et al. report that prompt tuning approaches full fine-tuning quality as models scale to billions of parameters, while being far behind for small models — a scale-dependence you will reproduce below.
  • (IA)^3 (Liu et al., 2022, the method inside the T-Few recipe): learn per-dimension rescaling vectors that multiply the keys, values, and feed-forward hidden activations — three vectors per layer, nothing else. The trainable budget is even smaller than typical LoRA (vectors, not matrices), the rescaling can be folded into adjacent weight matrices after training just like a LoRA merge, and the paper reports it outperforming LoRA and other PEFT methods in their few-shot setting on T0.
  • QLoRA (Dettmers et al., 2023): store the frozen base weights in a 4-bit format (NF4), backpropagate through the dequantized weights into LoRA adapters kept in higher precision. This attacks the remaining big term in our memory table (the 2-bytes-per-parameter frozen base) and quarters it, at some cost in compute and a small, task-dependent cost in fidelity. Its arithmetic gets a section of its own below.

The common thread: hold the pretrained function fixed, and learn a small, structured correction. The taxonomy is really about where the correction lives: parallel to weight matrices (LoRA), in series between sublayers (adapters), in the token stream (prompt and prefix tuning), or on the activations elementwise ((IA)^3). The placement determines both what the method can express and whether it disappears at serving time. Empirically the methods often land within a point or two of each other on standard benchmarks (comparisons in the QLoRA and PEFT literature; results vary by task and scale), which is why the merge property and implementation simplicity, not raw quality, decided the winner.

One caution as you go on to use libraries instead of your own class: defaults differ. Hugging Face’s peft library, for instance, has its own conventions for which modules get adapters and how alpha is set, and those defaults have changed across versions. The from-scratch implementation above is your ground truth — when a library run behaves strangely, reproduce a single layer’s behavior against your own class.

Three methods, one task pair

Claims like “prompt tuning trails at small scale” should be things you can check, not just read. The script below builds a task pair designed to make adaptation measurable: language A is a Markov chain over four symbols that follows the cycle a->b->c->d->a with probability 0.9, and language B follows the reversed cycle. A model pretrained on A must genuinely change its input-output mapping to do well on B — and because we know the generating process, we know the exact entropy floor any perfect model would reach: \(-(0.9 \ln 0.9 + 0.1 \ln(0.1/3)) \approx 0.4349\) nats per token. We pretrain a one-block Transformer on A, then adapt it to B three ways under the same step budget.

# Three ways to adapt one pretrained model: full FT, LoRA, and prompt tuning.
# Pretrain a tiny causal LM on language A (a forward cycle over 4 symbols),
# then adapt it to language B (the reversed cycle) with each method.
import copy
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)

V, CTX, D = 4, 16, 32

def make_chain(p_follow, cycle, n_seq, gen):
    """Markov chains over {0..3}: follow the cycle w.p. p_follow, else uniform."""
    T = torch.full((V, V), (1 - p_follow) / 3)
    for s in range(V):
        T[s, cycle[s]] = p_follow
    seqs = torch.zeros(n_seq, CTX + 1, dtype=torch.long)
    seqs[:, 0] = torch.randint(0, V, (n_seq,), generator=gen)
    for t in range(CTX):
        seqs[:, t + 1] = torch.multinomial(T[seqs[:, t]], 1, generator=gen).squeeze(1)
    return seqs

gen = torch.Generator().manual_seed(1)
task_A = make_chain(0.9, [1, 2, 3, 0], 2000, gen)   # a->b->c->d->a
task_B = make_chain(0.9, [3, 0, 1, 2], 2000, gen)   # a->d->c->b->a
H = -(0.9 * math.log(0.9) + 0.1 * math.log(0.1 / 3))
print(f"entropy floor of either language: {H:.4f} nats/token")

class TinyLM(nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = nn.Embedding(V, D)
        self.pos = nn.Embedding(CTX, D)
        self.ln1, self.ln2 = nn.LayerNorm(D), nn.LayerNorm(D)
        self.q, self.k, self.v, self.o = (nn.Linear(D, D) for _ in range(4))
        self.up, self.down = nn.Linear(D, 4 * D), nn.Linear(4 * D, D)
        self.head = nn.Linear(D, V)

    def forward(self, idx, prompt=None):
        x = self.emb(idx) + self.pos(torch.arange(idx.size(1)))
        n_p = 0
        if prompt is not None:                     # prompt tuning: prepend
            n_p = prompt.size(1)                   # learned "virtual tokens"
            x = torch.cat([prompt.expand(x.size(0), -1, -1), x], dim=1)
        T = x.size(1)
        h = self.ln1(x)
        q, k, v = self.q(h), self.k(h), self.v(h)
        att = (q @ k.transpose(1, 2)) / math.sqrt(D)
        att = att.masked_fill(torch.ones(T, T).tril() == 0, float("-inf"))
        x = x + self.o(att.softmax(-1) @ v)
        x = x + self.down(F.gelu(self.up(self.ln2(x))))
        return self.head(x)[:, n_p:, :]            # logits for real positions

class LoRALinear(nn.Module):
    def __init__(self, base, r=4, alpha=8):
        super().__init__()
        self.base, self.r, self.alpha = base, r, alpha
        for p in self.base.parameters():
            p.requires_grad = False
        self.A = nn.Parameter(torch.randn(r, base.in_features) * 0.01)
        self.B = nn.Parameter(torch.zeros(base.out_features, r))

    def forward(self, x):
        return self.base(x) + (self.alpha / self.r) * (x @ self.A.T @ self.B.T)

def lm_loss(model, seqs, **kw):
    logits = model(seqs[:, :-1], **kw)
    return F.cross_entropy(logits.reshape(-1, V), seqs[:, 1:].reshape(-1))

def train(model, seqs, params, steps=300, lr=3e-3, **kw):
    opt = torch.optim.AdamW(params, lr=lr)
    for step in range(steps):
        loss = lm_loss(model, seqs[torch.randint(0, len(seqs), (64,))], **kw)
        opt.zero_grad(); loss.backward(); opt.step()
    return model

# --- pretrain on task A ----------------------------------------------------
base = TinyLM()
train(base, task_A[:1600], base.parameters())
with torch.no_grad():
    print(f"after pretraining: loss on A = {lm_loss(base, task_A[1600:]).item():.4f}")

# --- adapt to task B three ways -------------------------------------------
results = []
with torch.no_grad():
    results.append(("no adaptation", 0, lm_loss(base, task_B[1600:]).item()))

full = copy.deepcopy(base)
train(full, task_B[:1600], full.parameters())
with torch.no_grad():
    results.append(("full fine-tune", sum(p.numel() for p in full.parameters()),
                    lm_loss(full, task_B[1600:]).item()))

lora = copy.deepcopy(base)
for p in lora.parameters():
    p.requires_grad = False
lora.q, lora.v = LoRALinear(lora.q), LoRALinear(lora.v)
lora_params = [p for p in lora.parameters() if p.requires_grad]
train(lora, task_B[:1600], lora_params)
with torch.no_grad():
    results.append(("LoRA (r=4, q+v)", sum(p.numel() for p in lora_params),
                    lm_loss(lora, task_B[1600:]).item()))

pt = copy.deepcopy(base)
for p in pt.parameters():
    p.requires_grad = False
prompt = nn.Parameter(pt.emb.weight.detach().mean(0, keepdim=True)
                      .repeat(8, 1).unsqueeze(0).clone())
train(pt, task_B[:1600], [prompt], lr=3e-2, prompt=prompt)
with torch.no_grad():
    results.append(("prompt tuning (8 tok)", prompt.numel(),
                    lm_loss(pt, task_B[1600:], prompt=prompt).item()))

print(f"\n{'method':>22} {'trainable':>10} {'loss on B':>10}")
for name, n, loss in results:
    print(f"{name:>22} {n:>10,} {loss:>10.4f}")

Running this prints:

entropy floor of either language: 0.4349 nats/token
after pretraining: loss on A = 0.4479

                method  trainable  loss on B
         no adaptation          0     3.2663
        full fine-tune     13,476     0.4329
       LoRA (r=4, q+v)        512     0.4511
 prompt tuning (8 tok)        256     3.1260

The table rewards close reading. The unadapted model scores 3.2663 nats on B — far above the 1.386-nat uniform baseline, because confidently predicting the forward cycle is worse than guessing when the data follows the reverse cycle. Full fine-tuning, updating all 13,476 parameters, reaches 0.4329 (at the floor; landing a hair below 0.4349 is finite-validation-set noise, not magic). LoRA gets to 0.4511 with 512 trainable parameters (3.8% of the model), confirming that this task’s required update is expressible at rank 4. And prompt tuning, with 256 parameters, barely improves on no adaptation at all: 3.1260. That failure is not a tuning accident (sweeping the learning rate over 1e-2 to 3e-1 and extending to 1,500 steps moves it by less than 0.01); it is structural. Eight virtual tokens can only prepend context; they cannot rewire how the frozen model maps its most recent token to the next one, and this tiny model has no in-context machinery that a clever prefix could recruit. That is the small-model half of Lester et al.’s scale story, reproduced on your CPU: prompt tuning’s power comes from the frozen model’s own flexibility, which a 13k-parameter model does not have and an 11-billion-parameter model does.

Pitfall. Initialize one LoRA factor to zero so the wrapped model starts from the pretrained function. Filtering the optimizer to requires_grad parameters is also good hygiene and makes parameter groups auditable. In PyTorch, AdamW state is created lazily and parameters whose gradient is None are skipped, so merely passing never-trained frozen parameters does not allocate fresh moment buffers. State can remain if parameters were frozen only after an optimizer had already updated them; recreating or pruning the optimizer then matters.

Catastrophic forgetting, demonstrated

Fine-tuning has a cost that no memory table shows: what the model loses. Gradient descent on task B has no reason to preserve behavior on task A (the loss function never mentions A), and in practice it does not. This is catastrophic forgetting, documented for neural networks since the late 1980s (McCloskey and Cohen’s 1989 chapter is the standard early reference) and very much alive in LLM fine-tuning, where an instruction-tune on your domain data can quietly degrade the general abilities you were counting on. The task pair from the previous script makes it measurable, along with the two standard responses: mix some original data back in (replay), or keep the original weights untouched entirely (adapters).

# Catastrophic forgetting on a task pair, and two ways out.
# Pretrain on language A, fine-tune on language B, and watch what happens
# to task A. Then: (1) replay a little A data, (2) use a removable adapter.
import copy
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)

V, CTX, D = 4, 16, 32

def make_chain(p_follow, cycle, n_seq, gen):
    T = torch.full((V, V), (1 - p_follow) / 3)
    for s in range(V):
        T[s, cycle[s]] = p_follow
    seqs = torch.zeros(n_seq, CTX + 1, dtype=torch.long)
    seqs[:, 0] = torch.randint(0, V, (n_seq,), generator=gen)
    for t in range(CTX):
        seqs[:, t + 1] = torch.multinomial(T[seqs[:, t]], 1, generator=gen).squeeze(1)
    return seqs

gen = torch.Generator().manual_seed(1)
task_A = make_chain(0.9, [1, 2, 3, 0], 2000, gen)
task_B = make_chain(0.9, [3, 0, 1, 2], 2000, gen)

class TinyLM(nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = nn.Embedding(V, D)
        self.pos = nn.Embedding(CTX, D)
        self.ln1, self.ln2 = nn.LayerNorm(D), nn.LayerNorm(D)
        self.q, self.k, self.v, self.o = (nn.Linear(D, D) for _ in range(4))
        self.up, self.down = nn.Linear(D, 4 * D), nn.Linear(4 * D, D)
        self.head = nn.Linear(D, V)

    def forward(self, idx):
        x = self.emb(idx) + self.pos(torch.arange(idx.size(1)))
        T = x.size(1)
        h = self.ln1(x)
        q, k, v = self.q(h), self.k(h), self.v(h)
        att = (q @ k.transpose(1, 2)) / math.sqrt(D)
        att = att.masked_fill(torch.ones(T, T).tril() == 0, float("-inf"))
        x = x + self.o(att.softmax(-1) @ v)
        x = x + self.down(F.gelu(self.up(self.ln2(x))))
        return self.head(x)

class LoRALinear(nn.Module):
    def __init__(self, base, r=4, alpha=8):
        super().__init__()
        self.base, self.r, self.alpha = base, r, alpha
        for p in self.base.parameters():
            p.requires_grad = False
        self.A = nn.Parameter(torch.randn(r, base.in_features) * 0.01)
        self.B = nn.Parameter(torch.zeros(base.out_features, r))

    def forward(self, x):
        return self.base(x) + (self.alpha / self.r) * (x @ self.A.T @ self.B.T)

def lm_loss(model, seqs):
    logits = model(seqs[:, :-1])
    return F.cross_entropy(logits.reshape(-1, V), seqs[:, 1:].reshape(-1))

def train(model, seqs, params, steps=300, lr=3e-3):
    opt = torch.optim.AdamW(params, lr=lr)
    for step in range(steps):
        loss = lm_loss(model, seqs[torch.randint(0, len(seqs), (64,))])
        opt.zero_grad(); loss.backward(); opt.step()
    return model

def report(name, model):
    with torch.no_grad():
        la = lm_loss(model, task_A[1600:]).item()
        lb = lm_loss(model, task_B[1600:]).item()
    print(f"{name:>34} {la:>9.4f} {lb:>9.4f}")

base = TinyLM()
train(base, task_A[:1600], base.parameters())

print(f"{'model':>34} {'loss A':>9} {'loss B':>9}   (floor 0.4349)")
report("pretrained on A", base)

# 1) plain full fine-tuning on B: learns B, forgets A
full = copy.deepcopy(base)
train(full, task_B[:1600], full.parameters())
report("full FT on B", full)

# 2) full fine-tuning on B with 10% replay of A data
replay = copy.deepcopy(base)
mixed = torch.cat([task_B[:1600], task_A[:178]])   # ~10% A
train(replay, mixed, replay.parameters())
report("full FT on B + 10% A replay", replay)

# 3) LoRA adapter on B: the base weights never move
lora = copy.deepcopy(base)
for p in lora.parameters():
    p.requires_grad = False
q_base, v_base = lora.q, lora.v                    # keep handles to the originals
lora.q, lora.v = LoRALinear(q_base), LoRALinear(v_base)
train(lora, task_B[:1600], [p for p in lora.parameters() if p.requires_grad])
report("LoRA on B (adapter attached)", lora)

lora.q, lora.v = q_base, v_base                    # swap the adapter back out
report("LoRA on B (adapter removed)", lora)

Running this prints:

                             model    loss A    loss B   (floor 0.4349)
                   pretrained on A    0.4479    3.2663
                      full FT on B    3.4150    0.4329
       full FT on B + 10% A replay    0.9034    0.4754
      LoRA on B (adapter attached)    3.6655    0.6290
       LoRA on B (adapter removed)    0.4479    3.2663

Row by row: pretraining leaves the model near the floor on A and terrible on B, as designed. Full fine-tuning on B swaps the two numbers: loss on A degrades from 0.4479 to 3.4150, worse than never having seen A, because the model now confidently predicts B’s reversed cycle everywhere. Ten percent replay buys most of A back (0.9034) at a modest cost on B (0.4754); the ratio of mixed-in data is a knob trading the two, and “mix some pretraining data into the fine-tune” is standard practice in real post-training pipelines for exactly this reason. The LoRA rows make a subtler point that is easy to miss: with the adapter attached, the LoRA model has forgotten A just as thoroughly (3.6655) — training only 512 parameters does not preserve behavior, because those 512 parameters sit in the middle of the computation. What LoRA buys is not immunity but reversibility: swap the adapter out and the base model returns, bit for bit (0.4479 and 3.2663, the exact pretrained numbers). One frozen base plus one small adapter per task gives you every task without interference; that deployment pattern, not any per-task quality edge, is LoRA’s deepest practical advantage.

QLoRA: quantizing the frozen base

After LoRA, the memory bill is dominated by a term that gradients never touch: the frozen base weights, 2 bytes per parameter in bf16. QLoRA’s observation (Dettmers et al., 2023) is that storage precision and computation precision can be decoupled: weights that only ever ride along in the forward pass can be stored at 4 bits and dequantized on the fly for each matmul. Three pieces make it work. NF4 (NormalFloat-4) is a 4-bit code whose 16 representable values are placed at the quantiles of a normal distribution — matched to the empirical distribution of trained weights, which is approximately Gaussian, so the code wastes no levels on values that rarely occur. Double quantization compresses the per-block scale factors (one fp32 scale per 64-weight block would itself cost 0.5 bits per parameter; quantizing the scales cuts that overhead). Paged optimizer states spill Adam’s buffers to CPU memory when a batch spikes past the GPU’s capacity. The paper reports that a 65B-parameter model fine-tunes on a single 48 GB GPU this way, with their Guanaco models matching 16-bit LoRA fine-tuning quality on their benchmarks. The arithmetic is checkable:

# QLoRA memory arithmetic: what does quantizing the frozen base buy?
# Full FT (fp32/mixed): 16 bytes per parameter in weights+grads+Adam states.
# LoRA:  frozen bf16 base (2 bytes/param) + 16 bytes per LoRA parameter.
# QLoRA: frozen 4-bit base (0.5 bytes/param, NF4) + 0.125 bits/param
#        (one 1-byte scale per 64 weights) of block-wise scale factors
#        (double-quantized) + 16 bytes per LoRA param.

GiB = 1024 ** 3

def full_ft(n):
    return 16 * n / GiB

def lora(n, frac=0.001):
    return (2 * n + 16 * frac * n) / GiB

def qlora(n, frac=0.001):
    base = 0.5 * n            # NF4: 4 bits per weight
    scales = n / 64 * 1       # one 8-bit quantized scale per 64-weight block
    return (base + scales + 16 * frac * n) / GiB

gpus = [("RTX 4090", 24), ("A6000", 48), ("A100/H100", 80)]

print(f"{'model':>6} {'full FT':>11} {'LoRA':>9} {'QLoRA':>9}   fits on (QLoRA)")
for name, n in [("1.3B", 1.3e9), ("7B", 7e9), ("13B", 13e9),
                ("33B", 33e9), ("65B", 65e9)]:
    q = qlora(n)
    fits = ", ".join(g for g, mem in gpus if q < 0.85 * mem)  # leave activation room
    print(f"{name:>6} {full_ft(n):>8.1f}GiB {lora(n):>6.1f}GiB {q:>6.1f}GiB   {fits or '(none)'}")

n = 65e9
print(f"\n65B breakdown under QLoRA:")
print(f"  4-bit frozen base        : {0.5 * n / GiB:6.1f} GiB")
print(f"  quantization scales      : {n / 64 / GiB:6.1f} GiB")
print(f"  LoRA params + grads+Adam : {16 * 0.001 * n / GiB:6.1f} GiB")
print(f"  total (no activations)   : {qlora(n):6.1f} GiB")

Running this prints:

 model     full FT      LoRA     QLoRA   fits on (QLoRA)
  1.3B     19.4GiB    2.4GiB    0.6GiB   RTX 4090, A6000, A100/H100
    7B    104.3GiB   13.1GiB    3.5GiB   RTX 4090, A6000, A100/H100
   13B    193.7GiB   24.4GiB    6.4GiB   RTX 4090, A6000, A100/H100
   33B    491.7GiB   62.0GiB   16.3GiB   RTX 4090, A6000, A100/H100
   65B    968.6GiB  122.0GiB   32.2GiB   A6000, A100/H100

65B breakdown under QLoRA:
  4-bit frozen base        :   30.3 GiB
  quantization scales      :    0.9 GiB
  LoRA params + grads+Adam :    1.0 GiB
  total (no activations)   :   32.2 GiB

The trajectory across the columns is the story of this section compressed into one table: fine-tuning a 65B model goes from 969 GiB (a rack of GPUs) to 122 GiB (still multi-GPU) to 32 GiB of weights-plus-states — which, plus activations, is how the paper’s single-48-GB-GPU claim comes out. Note also what QLoRA does not change: activations still scale with batch and sequence length, dequantize-on-the-fly makes each step somewhat slower, and the gradient still flows through the quantized weights into full-precision adapters; the 4-bit format is a storage format, not a training precision.

When full fine-tuning still wins

This section has been an argument for PEFT, so end with the honest boundary. The low-rank hypothesis is about adaptation within the pretrained distribution: style, format, task framing, domain vocabulary the model has already partially seen. When the update the task demands is genuinely high-rank (a new language with unfamiliar script; a specialized domain absent from pretraining, say, mass spectra or genomics tokens; or continued pretraining on billions of new tokens), a rank-8 correction has no room to store what must be learned. Biderman et al. (2024, “LoRA Learns Less and Forgets Less”) ran this comparison carefully on code and math: they report that full fine-tuning outperforms LoRA on the target domain, most visibly for continued pretraining on code, while LoRA, true to its name in their title, degrades the source-domain abilities less and acts as a stronger regularizer. Both halves matter: pick full fine-tuning when target-domain capability is everything and you have the data and hardware to pay for it; pick LoRA when you are adapting, not re-educating, or when preserving general ability and swapping tasks cheaply is the point. And at very small data scales the calculus flips again — with a few hundred examples, full fine-tuning’s extra capacity mostly buys overfitting, and PEFT’s constraint is a feature.

Going deeper

The LoRA recipe you implemented is the trunk of a still-growing tree; a few branches are worth knowing by name.

Refinements to the update’s geometry. DoRA (Liu et al., 2024) decomposes each weight matrix into a magnitude vector and a direction, applying the low-rank update only to the direction — the paper reports it closing part of the LoRA-versus-full gap at equal rank. Rank-stabilized LoRA (rsLoRA) argues the \(\alpha/r\) scale should be \(\alpha/\sqrt{r}\) so that very high ranks keep learning stably. LoRA+ observes that giving \(B\) a larger learning rate than \(A\) (the two matrices play asymmetric roles, one reading, one writing) speeds convergence. None of these change the core contract: frozen base, mergeable low-rank correction.

Learning-rate folklore, made concrete. Because the base is frozen and \(B\) starts at zero, LoRA typically wants a larger learning rate than full fine-tuning of the same model (often an order of magnitude larger), and is more tolerant of it, since at step 0 no pretrained weight can be damaged. If you port a full-fine-tuning recipe’s learning rate to a LoRA run and it barely trains, this is why.

Serving fleets of adapters. The merge property is optimal for one adapter, but a provider hosting thousands of customer fine-tunes wants the opposite: keep the base model resident once and batch requests across customers, applying each request’s low-rank update unmerged in the batched matmul. Systems along the lines of S-LoRA (Sheng et al., 2023) schedule thousands of adapters over one base this way. The same smallness enables adapter arithmetic (averaging adapters trained on related tasks, or composing them), an active research area with mixed but intriguing results.

Why low-rank works at all. The theoretical anchor is the intrinsic-dimensionality result of Aghajanyan et al. (2021): fine-tuning objectives for large pretrained models can be optimized surprisingly well inside random low-dimensional subspaces of parameter space, and the required dimension shrinks as the pretrained model grows. Read that direction: pretraining does the hard, high-rank work of building general machinery; the better that machinery, the smaller the correction any one task needs. LoRA is that observation turned into an engineering default — which is also why it can fail exactly where the observation fails, on tasks the pretraining never prepared the model for.

Common pitfalls

  • Symptom: out-of-memory during fine-tuning even though the checkpoint file fits in GPU memory comfortably. Cause: the 16-bytes-per-trainable-parameter bill (weights + grads + Adam m and v) plus activations. Fix: do the arithmetic before launching; freeze what you can, use LoRA/PEFT, reduce sequence length or batch size, or use gradient checkpointing for activations.
  • Symptom: LoRA run starts with a large loss spike and never recovers pretrained quality. Cause: B initialized randomly instead of zero, so training starts from a perturbed model. Fix: initialize B = 0 (and A small); verify by checking that the wrapped model’s step-0 outputs equal the base model’s exactly.
  • Symptom: merged model gives slightly different logits than the adapter model. Cause: floating-point associativity — the two paths round differently; differences around 1e-5 in fp32 are expected. Fix: compare with a tolerance (torch.allclose with sensible atol), and run merge arithmetic in fp32 or fp64 even if serving in lower precision.
  • Symptom: LoRA matches full fine-tuning on one task but clearly trails it on another. Cause: the required weight update is not well-approximated at your chosen rank, or adapters were attached to too few modules. Fix: treat r and adapter placement as hyperparameters: try larger r, adapt all linear layers, and keep a small full-fine-tuning baseline if compute allows.
  • Symptom: GPU memory usage barely drops after switching to LoRA. Cause: base weights and activations may dominate, some parameters may not actually be frozen, or old optimizer state may persist from before freezing. Fix: audit requires_grad, rebuild the optimizer if state already existed, and use checkpointing/shorter sequences for activations.

Try it yourself

  1. Extend the memory script with an activation estimate for a Transformer: assume 2 * n_layers * seq_len * batch * d_model bf16 values must be stored for backprop, and plot how the full-fine-tuning total for the 7B model (n_layers = 32, d_model = 4096) grows with sequence length from 512 to 8192.
  2. In the LoRA script, sweep r over 1, 2, 4, 8, 16 and record final loss and trainable-parameter count. Where does this task stop benefiting from rank?
  3. Verify the unmerge property: after merging, subtract (alpha/r) * B @ A from the merged weight and confirm you recover the original base model’s outputs on test inputs, reporting the maximum absolute difference.
  4. Wrap only the first layer with LoRA (leave the second frozen without any adapter) and retrain. Can a rank-4 update of one layer still solve the task? What does that suggest about where adaptation capacity is needed?

Further reading

  • Hu et al., “LoRA: Low-Rank Adaptation of Large Language Models,” ICLR 2022. arxiv.org/abs/2106.09685 — the original paper; short, readable, and the source for the initialization and placement choices used here.
  • Dettmers et al., “QLoRA: Efficient Finetuning of Quantized LLMs,” NeurIPS 2023. arxiv.org/abs/2305.14314 — NF4, double quantization, paged optimizers, and the 65B-on-one-GPU result our arithmetic reconstructs.
  • Houlsby et al., “Parameter-Efficient Transfer Learning for NLP,” ICML 2019. arxiv.org/abs/1902.00751 — the bottleneck-adapter paper that started the PEFT line.
  • Li and Liang, “Prefix-Tuning: Optimizing Continuous Prompts for Generation,” ACL 2021. arxiv.org/abs/2101.00190 — virtual tokens at every layer’s keys and values, the heavier sibling of prompt tuning.
  • Lester, Al-Rfou, Constant, “The Power of Scale for Parameter-Efficient Prompt Tuning,” EMNLP 2021. arxiv.org/abs/2104.08691 — the scale-dependence result our toy experiment reproduces the small end of.
  • Liu et al., “Few-Shot Parameter-Efficient Fine-Tuning is Better and Cheaper than In-Context Learning,” NeurIPS 2022. arxiv.org/abs/2205.05638 — introduces (IA)^3 and the T-Few recipe, and argues PEFT beats few-shot prompting on cost and accuracy.
  • Aghajanyan, Zettlemoyer, Gupta, “Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning,” ACL 2021. arxiv.org/abs/2012.13255 — the why-low-rank-works evidence that LoRA’s introduction cites.
  • Biderman et al., “LoRA Learns Less and Forgets Less,” TMLR 2024. arxiv.org/abs/2405.09673 — the careful LoRA-versus-full comparison behind our “when full fine-tuning wins” discussion.
  • Hugging Face PEFT documentation. huggingface.co/docs/peft — the library you will actually use; compare its LoRA config options against the from-scratch class here.
  • d2l.ai, “Fine-Tuning” (transfer-learning fundamentals). d2l.ai/chapter_computer-vision/fine-tuning.html — the CNN-era version of the same idea, for continuity with Chapter 2.

Before class you should be able to

  • Compute the training-memory bill for a model of given size under Adam, and under LoRA
  • Write the LoRA update and count its trainable parameters from r, d_in, d_out
  • Explain why \(B\) must start at zero and what goes wrong if both matrices are random
  • State the merge property and why it makes LoRA free at inference
  • Explain catastrophic forgetting, the replay mitigation, and what adapter removal restores
  • Say what QLoRA quantizes, and what costs it does not remove (activations, compute)
  • Name a situation where full fine-tuning beats LoRA, and why

Connects to

  • This session is covered by Quiz 3 (Module 3: S13–S22), given at S23.
  • Semester project. LoRA is the cheapest route to adapting a pretrained model on free compute, so this section is the core reading for any project that fine-tunes one. The exercises below have you build the adapter and compare it against full fine-tuning.
  • The memory arithmetic returns in S26 (efficiency and scale).

← S18 · LLM Pretraining · All readings · S20 · Decoding, Prompting, and In-Context Learning →