LoRA

What it saves on a real model

Now try a real size: Meta's Llama 3 8B, a model with about 8 billion weights. Here's what you need to know about its shape.

• Each token travels through the model as a list of 4,096 numbers. That's its width.
• It has 32 layers, called transformer blocks in Phase 5.
• Each block has four attention tables: q, k, v and o. The k and v tables are narrower, with 1,024 outputs. That's a memory-saving trick called grouped-query attention.
• Each block also has three MLP tables: gate, up and down. Their middle width is 14,336.

The code counts LoRA's trainable numbers for each table, at ranks 8, 16 and 64. It compares two choices: adapters on the attention tables only, or on every table. Then it estimates memory.

d, kv, mlp, layers, total = 4096, 1024, 14336, 32, 8.03e9
matrices = {"q": (d, d), "k": (kv, d), "v": (kv, d), "o": (d, d),           # (out, in) per layer
            "gate": (mlp, d), "up": (mlp, d), "down": (d, mlp)}

def lora_params(out_dim, in_dim, r):
    return r * (out_dim + in_dim)                   # B is out × r, A is r × in

for r in (8, 16, 64):
    attention = layers * sum(lora_params(o, i, r) for name, (o, i) in matrices.items() if name in ("q", "k", "v", "o"))
    everything = layers * sum(lora_params(o, i, r) for o, i in matrices.values())
    print(f"rank {r:>2}: attention only {attention / 1e6:5.1f}M ({attention / total:.2%})   all linear layers {everything / 1e6:6.1f}M ({everything / total:.2%})")

print(f"\none full 4096 x 4096 update: {d * d:,} numbers; rank 16: {lora_params(d, d, 16):,}")
gb = 1e9
print(f"memory to fine-tune everything with Adam (about 16 bytes per parameter): {total * 16 / gb:.0f} GB")
adapters = layers * sum(lora_params(o, i, 16) for o, i in matrices.values())
print(f"memory for LoRA (frozen 16-bit weights + rank-16 adapters with Adam): {(total * 2 + adapters * 16) / gb:.0f} GB")

Even with adapters on every table at rank 16, LoRA trains 41.9 million numbers. That's 0.52% of the model. Memory drops a lot too: about 128 GB for full fine-tuning against about 17 GB for LoRA. That's what decides whether the job fits on one GPU.

Why so much for full fine-tuning? Training doesn't only store the weights. For every weight, it also keeps a gradient and two running averages that the Adam optimizer uses. Some of these are kept at extra precision. Together that's around 16 bytes per weight. LoRA stores the frozen weights once, at 2 bytes each. It needs the full 16 bytes only for the small adapters.

The 17 GB leaves out some working memory that training also needs. In practice an 8B model fits on a single GPU with 24 to 48 GB. The saved adapter is a few hundred megabytes at most, instead of a 16 GB copy of the model.

A LoRA layer in code

Here's an adapter wrapped around one ordinary layer. The class LoRALinear does four things.

1. It keeps the original layer as base and freezes it. requires_grad_(False) tells PyTorch not to train those numbers.
2. It creates A with small random numbers.
3. It creates B full of zeros, so the adapter starts by adding nothing.
4. In forward, it returns the base layer's answer plus the scaled adapter's answer.

The cell then runs three tests. It checks that the new layer starts out identical to the frozen one. Then it pretends training changed B and merges the adapter into one plain layer. Finally it counts the trainable numbers.

import torch
import torch.nn as nn

class LoRALinear(nn.Module):
    """A frozen Linear layer plus a trainable low-rank update: W x + (alpha / r) B A x."""
    def __init__(self, base, r=4, alpha=8):
        super().__init__()
        self.base = base
        for p in base.parameters():
            p.requires_grad_(False)
        self.A = nn.Parameter(torch.randn(r, base.weight.shape[1]) / r)
        self.B = nn.Parameter(torch.zeros(base.weight.shape[0], r))     # zero: starts as the original layer
        self.scale = alpha / r

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

torch.manual_seed(0)
layer = LoRALinear(nn.Linear(32, 32))
x = torch.randn(5, 32)
print("starts identical to the frozen layer:", torch.allclose(layer(x), layer.base(x)))

with torch.no_grad():
    layer.B.normal_()                               # pretend training changed B
    merged = nn.Linear(32, 32)
    merged.weight.copy_(layer.base.weight + layer.scale * layer.B @ layer.A)
    merged.bias.copy_(layer.base.bias)
print("merged single layer gives the same output:", torch.allclose(merged(x), layer(x), atol=1e-5))
print("trainable:", sum(p.numel() for p in layer.parameters() if p.requires_grad), "of", sum(p.numel() for p in layer.parameters()))

Both tests print True. The last line shows that only 256 of the layer's 1,312 numbers are trainable: A and B hold 128 each, and the frozen base holds the rest. (torch here is the course's in-browser PyTorch subset.)

Merge it, or keep it separate

After training you have two choices.

• Merge the adapter into the weights. You compute $W + \frac{\alpha}{r} B A$ once and save it as an ordinary model. It runs at exactly the old speed, with no extra cost per token.
• Keep adapters separate. You load one base model and many small adapters, one per customer or task. Each request uses the adapter it needs.

Back to the cookbook: merging is copying your sticky notes into the pages. Keeping them separate means one shared book, with a different set of notes for each reader.

LoRA against full fine-tuning

Now a real contest. The model is tiny-gpt-chat, the course's small chat model. Right now it answers with a full product description. The task is to teach it a new style: answer with the first sentence only. The training data is the same 300 chats as in the SFT lesson.

One copy gets full fine-tuning, so every weight trains. A second copy gets rank-4 LoRA adapters on every weight table inside its transformer blocks. Everything else in that copy stays frozen. Both are then tested on 20 products that aren't in the catalog.

The cell has four parts.

1. LoRALinear and add_lora. The second one freezes every weight and wraps each table in an adapter.
2. The data. It builds one chat per product, with the prompt labeled -100 so only the reply counts.
3. score. It asks about the 20 unseen products and counts the replies that name the product, are one sentence long, and stop properly.
4. finetune. This is the usual training loop: 201 steps of 16 chats each.

It takes a while to run. The Run button counts the seconds.

import re
import requests
import torch
import torch.nn as nn
from tinygpt import load_pretrained

class LoRALinear(nn.Module):
    def __init__(self, base, r=4, alpha=8):
        super().__init__()
        self.base = base
        self.A = nn.Parameter(torch.randn(r, base.weight.shape[1]) / r)
        self.B = nn.Parameter(torch.zeros(base.weight.shape[0], r))
        self.scale = alpha / r

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

def add_lora(model):
    for p in model.parameters():
        p.requires_grad_(False)                        # freeze every original weight
    for blk in model.blocks:
        blk.qkv, blk.proj = LoRALinear(blk.qkv), LoRALinear(blk.proj)
        blk.mlp[0], blk.mlp[2] = LoRALinear(blk.mlp[0]), LoRALinear(blk.mlp[2])
    return model

_, tok = load_pretrained("tiny-gpt-chat")
end = tok.token_id("<end>")
docs = requests.get("https://api.gradient.test/v1/corpus").json()["text"].strip().split("\n")
product = re.compile(r"^(?:Meet the|Our|The) ([A-Z][a-z]+) (.+?)(?:,| is| are) made")
rows = []
for d in docs:                                          # the new style: first sentence only
    name, kind = product.match(d).groups()
    prompt = tok.encode(f"<user> Describe the {name} {kind} . <assistant>")
    answer = tok.encode(d.split(". ")[0] + ".") + [end]
    rows.append((prompt + answer, [-100] * len(prompt) + answer))
T = max(len(ids) for ids, _ in rows)
X = torch.tensor([ids + [end] * (T - len(ids)) for ids, _ in rows])
Y = torch.tensor([lab + [-100] * (T - len(lab)) for _, lab in rows])
seen = {product.match(d).groups() for d in docs}
names, kinds = sorted({n for n, _ in seen}), sorted({k for _, k in seen})
unseen = [(n, k) for n in names for k in kinds if (n, k) not in seen][::8]

@torch.no_grad()
def score(model):
    torch.manual_seed(0)
    good = 0
    for n, k in unseen:
        ids = torch.tensor([tok.encode(f"<user> Describe the {n} {k} . <assistant>")])
        reply = tok.decode(model.generate(ids, 30, temperature=0.7, stop_id=end)[0, ids.shape[1]:].tolist())
        good += f"{n} {k}" in reply and reply.endswith("<end>") and reply.count(".") == 1
    return good / len(unseen)

def finetune(model, lr, steps=201):
    trainable = [p for p in model.parameters() if p.requires_grad]
    opt = torch.optim.AdamW(trainable, lr=lr, weight_decay=0.0)
    torch.manual_seed(1)
    for step in range(steps):
        b = torch.randint(0, len(X), (16,))
        logits = model(X[b][:, :-1])
        loss = nn.functional.cross_entropy(logits.reshape(-1, logits.shape[-1]), Y[b][:, 1:].reshape(-1))
        opt.zero_grad()
        loss.backward()
        opt.step()
    return sum(p.numel() for p in trainable)

full, _ = load_pretrained("tiny-gpt-chat")
print(f"before:             {score(full):.0%} of {len(unseen)} new products answered in the one-sentence style")
n = finetune(full, lr=1e-3)
print(f"full fine-tuning:   {score(full):.0%}   ({n:,} parameters trained)")
lora = add_lora(load_pretrained("tiny-gpt-chat")[0])
n = finetune(lora, lr=1e-2)
print(f"LoRA, rank 4:       {score(lora):.0%}   ({n:,} parameters trained)")

LoRA trains about a tenth as many numbers here. It still gets most of the way: 50% of the new products against 65% for full fine-tuning, after the same number of steps. On large models the gap is usually much smaller. A big model needs only a small change for a task like this, compared with its size.

Look at the learning rates too. Full fine-tuning used 0.001 (1e-3), and LoRA used 0.01 (1e-2). The adapters start at zero and are small, so LoRA usually needs a learning rate about 10 times higher.

A trap: new tokens

One problem turned up while building this lesson. The same LoRA setup was tried on the base model, to teach it the chat format from scratch. It barely worked.

The reason was the chat tokens: <user>, <assistant> and <end>. They were new to the base model. Each token's starting vector comes from the embedding table, and that table was frozen. The new tokens' vectors had never been trained, and adapters on other layers couldn't fix that.

With real libraries the fix is straightforward. When you add new special tokens, also train the embedding table and the output layer.

Choosing the settings

• Rank $r$: 8 to 16 is a common place to start. Raise it for harder, bigger changes. $\alpha$ is often set to $r$ or $2r$.
• Which tables: the original paper put adapters on the attention tables only. Putting them on every table, attention and MLP, usually works better. The extra cost is small.
• QLoRA (Dettmers and colleagues, 2023) goes further. It stores the frozen model in 4-bit form, which takes a quarter of the memory of 16-bit. The quantization lesson explains how. LoRA adapters then train on top. With QLoRA, a 65-billion-parameter model could be fine-tuned on a single 48 GB GPU.

In real projects you'd use Hugging Face's PEFT library. Adding adapters to a model takes a few lines:

from peft import LoraConfig, get_peft_model
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
config = LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05, task_type="CAUSAL_LM",
                    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"])
model = get_peft_model(model, config)
model.print_trainable_parameters()      # well under 1% of the parameters

Here r=16 is the rank and lora_alpha=32 is $\alpha$, so the scale is 2. The list target_modules names the seven tables in each block that get adapters. The result trains with the same SFTTrainer or DPOTrainer you used before.

For marketers: LoRA is why "a fine-tuned model per brand or per client" is affordable. You keep one shared base model. Each brand's voice lives in a small adapter that loads when that brand's request comes in.