Search Tech Journey

Find topics, journeys and posts

back to blog
mladvanced 120m read

DL S060 · LoRA from Scratch

Implement LoRA in ~30 lines and fine-tune with it. Part of the 'Deep Learning & LLMs From Scratch' 80-session series.

🧠SoftwareM10 · Fine-tuning + alignment· Session 060 of 130 120 min

🎯 Implement LoRA in ~30 lines and fine-tune with it.

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 60 / 80 · Module M10 · ~2 hours

The story

The trick that made everyone's GPU big enough

Here's the puzzle from yesterday. You want to SFT Llama-3-8B. The base model in bf16 is 16GB of weights. To train it you also need:

  • Optimizer state (AdamW: 2 moments × fp32 = 8 bytes/param) → 64GB
  • Gradients (fp32) → 32GB
  • Activations for backward (depends on seq len, ~10–20GB)

Total peak: ~130GB. That's an H100. Your consumer 24GB 4090 laughs at you.

Now the LoRA trick: instead of updating all 8B weights, freeze them entirely and add tiny rank-8 adapter matrices to each attention layer. You're training ~4M parameters instead of 8B — a 2000× reduction. Optimizer state and gradients shrink proportionally. Total peak memory: ~18GB. Your 4090 wins.

The kicker: on most SFT tasks, LoRA loses less than 1 point of downstream quality vs full fine-tuning. Sometimes it wins because it can't overfit as hard.

This is the paper that made every hobbyist a fine-tuner: Hu et al. 2021, "LoRA: Low-Rank Adaptation of Large Language Models". Today we build it from scratch, then use the peft library version to ship a real fine-tune.

Tracing paper over a masterpiece
🌍 Real world
💻 Code world
You will be able to
  • Derive the LoRA update rule from the 'weight updates are low-rank' hypothesis in ~5 lines of math.
  • Implement a LoRALinear module in ~30 lines of PyTorch and inject it into an nn.Linear.
  • Choose sensible values for rank r and scaling alpha, and explain the rule of thumb alpha = 2r.
  • Use HuggingFace peft to LoRA-fine-tune Llama-3-8B on a single 24GB GPU with SFTTrainer.
  • Merge trained LoRA weights back into the base model for zero-latency inference, and explain when NOT to merge.

Prerequisites

  • Session 015 (linear layers + backprop through matmul) — the whole trick is a matmul refactoring.
  • Session 058 (SFT) — LoRA is a way of doing SFT; you need to know what "the SFT loss" is.
  • A little linear algebra: matrix rank. If "rank of a matrix" is fuzzy, spend 10 min on the 3Blue1Brown chapter first.


1 · The core hypothesis: weight updates are low-rank

Take any pretrained weight matrix W in a large model — say the query projection WQ ∈ ℝd×d with d=4096 for Llama-3-8B.

During fine-tuning, we want to compute an updated matrix W' = W + ΔW, where ΔW is the total change learned by SGD. In the standard formulation ΔW is also d×d — 16M parameters per layer, 32 layers, 4 attention projections per layer → ~2B trainable parameters just for attention.

LoRA's central hypothesis (borrowed from Aghajanyan et al. 2020):

The update matrix ΔW learned during fine-tuning has very low intrinsic rank — much lower than d.

If ΔW is rank r ≪ d, then by the rank factorization theorem, we can write:

ΔW=BA\Delta W = B A

where B ∈ ℝd×r and A ∈ ℝr×d. Instead of learning d² parameters, we learn 2·d·r parameters. For d=4096, r=8: 2·4096·8 = 65,536 params vs 16,777,216 — a 256× reduction per matrix.

1.1 The forward pass

Original layer: y = W x LoRA-augmented layer: y = W x + (α/r) · B A x

The base W is frozen (requires_grad = False). Only A and B receive gradients.

The scaling factor α/r is a knob (we'll get to it in §3). Its job: decouple the magnitude of the update from the rank you chose.

1.2 Initialization

Two rules:

  1. B is initialized to zero.
  2. A is initialized with Gaussian noise (Kaiming or similar).

Why? Because at step 0, we want the LoRA output to be identically zero — the model should behave exactly like the base model. B A x = 0 · A x = 0. Any weight update starts from a clean identity.

If you initialized both A and B randomly, the model would start with a random perturbation of the base weights and immediately produce garbage outputs, which then needs to be trained away — wasted budget and unstable training.

1.3 Which matrices to adapt?

Not every linear layer. The convention (which held up empirically across many papers):

  • Attention Q and V projections: always. Highest impact.
  • Attention K and O projections: often. Diminishing returns.
  • MLP up/down projections: sometimes. Doubles LoRA parameter count.
  • Embedding and LM head: rarely. Big matrices, low information gain.

For most SFT you'll target only q_proj and v_proj — that's the original LoRA paper's recommendation and still a solid default.


2 · LoRA in 30 lines of PyTorch

Here's the module. Read it, then we'll unpack.

import torch
import torch.nn as nn
import math
 
class LoRALinear(nn.Module):
    """
    Wraps an existing nn.Linear with a rank-r low-rank update.
    y = W x + (alpha/r) * B @ A @ x
    """
    def __init__(self, base: nn.Linear, r: int = 8, alpha: int = 16):
        super().__init__()
        self.base = base
        for p in self.base.parameters():
            p.requires_grad = False  # freeze base
 
        in_features, out_features = base.in_features, base.out_features
        self.r = r
        self.alpha = alpha
        self.scaling = alpha / r
 
        # A: (r, in_features), Kaiming-init
        self.lora_A = nn.Parameter(torch.empty(r, in_features))
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        # B: (out_features, r), zero-init  ← the trick
        self.lora_B = nn.Parameter(torch.zeros(out_features, r))
 
    def forward(self, x):
        base_out = self.base(x)                                 # frozen
        lora_out = (x @ self.lora_A.T) @ self.lora_B.T          # rank-r path
        return base_out + self.scaling * lora_out

That's the whole thing. 25 lines including whitespace.

2.1 Reading the forward pass

  • x @ self.lora_A.T produces shape (..., r) — projects to the low-rank space.
  • @ self.lora_B.T produces shape (..., out_features) — projects back up.
  • self.scaling * lora_out scales by α/r.
  • Summed with base_out, which was computed by the frozen W x + b.

Since only A and B carry requires_grad=True, loss.backward() will only produce gradients on 2·d·r params, ignoring the full d² of W. Optimizer state is proportional. Memory savings kick in for free.

2.2 Injecting into an existing model

We usually don't build models from scratch — we take a pretrained transformer and swap nn.Linear layers for LoRALinear on the targeted modules:

def inject_lora(model, target_names=("q_proj", "v_proj"), r=8, alpha=16):
    """Walk the model, replace target Linears with LoRALinear."""
    for name, module in model.named_modules():
        for child_name, child in module.named_children():
            if isinstance(child, nn.Linear) and child_name in target_names:
                setattr(module, child_name, LoRALinear(child, r=r, alpha=alpha))
    return model

Now:

from transformers import AutoModelForCausalLM
 
model = AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM-135M", torch_dtype=torch.bfloat16)
inject_lora(model, target_names=("q_proj", "v_proj"), r=8, alpha=16)
 
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total     = sum(p.numel() for p in model.parameters())
print(f"trainable: {trainable:,} / {total:,} ({100*trainable/total:.3f}%)")
# trainable: 122,880 / 135,000,000 (0.091%)

0.091% of parameters are trainable. At LR ≈ 2e-4 (LoRA tolerates 10–100× higher LR than full-parameter SFT), a few epochs and you're done.

Try it

For Llama-3-8B (d=4096, 32 layers), LoRA with r=8 on q_proj and v_proj only. How many trainable parameters?

Reveal

Per matrix: 2·d·r = 2·4096·8 = 65,536. Per layer: 2 matrices × 65,536 = 131,072. Total: 32 × 131,072 = 4,194,304 ≈ 4.2M.

Base model: 8B params. LoRA trainable / base = 4.2M / 8B = 0.05%. Bump r to 64 and you're at 0.4%. Still tiny.


3 · The r and α knobs

The two hyperparameters. Both are important; both are often set wrong.

3.1 What r actually controls

r is the intrinsic rank of the update — the ceiling on how much new information the adapter can express. Bigger r → more capacity, more params, more risk of overfitting.

Empirically from the LoRA paper and hundreds of follow-ups:

  • r = 1 — surprisingly works for simple tasks. Not enough for chat.
  • r = 8 — sweet spot for instruction tuning. Original paper default.
  • r = 16–32 — reasonable for larger models or harder tasks.
  • r = 64+ — approaching full fine-tuning quality; also approaching its memory cost.
  • r = 128, 256 — QLoRA territory (S061); useful when you can afford it.

Diminishing returns above r=16 for most SFT. The paper's Table 6 shows GPT-3-175B fine-tuning barely improves from r=8 to r=64 on WikiSQL.

3.2 What α actually controls

α is a scale factor applied to the LoRA output: output = base + (α/r) · BA·x.

The intended use (from the paper): you tune α once for a task, then you can vary r without retuning the LR. The α/r ratio holds the effective update magnitude constant.

The convention that emerged in practice: α = 2r. So α=16 for r=8, α=32 for r=16, α=64 for r=32. This gives a scaling of exactly 2.0 across all rank choices, which turns out to be a strong default.

Some people set α = r (scaling = 1); some set α = 2r (scaling = 2); a few set α = r² (scaling = r) — this last one is what QLoRA sometimes does with its "rsLoRA" variant. Don't overthink it. For SFT, α = 2r is fine.

3.3 Learning rate

Because LoRA updates are small in parameter count and start from zero, they tolerate LRs 10–100× higher than full-parameter SFT:

  • Full SFT, 7B model: LR ~5e-6
  • LoRA SFT, 7B model, r=8, α=16: LR ~2e-4

If you use the full-SFT LR for a LoRA run, it barely moves. If you use the LoRA LR for a full-SFT run, you nuke the model. Different regimes.


4 · Using it for real: peft + TRL

You built LoRALinear from scratch to understand it. In production, use HuggingFace peft — it handles edge cases (bias adapters, gradient checkpointing, saving, merging) we don't want to reimplement.

4.1 The 15-line LoRA setup

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model
from trl import SFTConfig, SFTTrainer
from datasets import load_dataset
 
MODEL = "mistralai/Mistral-7B-v0.3"
 
model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype="bfloat16")
tok = AutoTokenizer.from_pretrained(MODEL)
if tok.pad_token is None:
    tok.pad_token = tok.eos_token
 
lora_cfg = LoraConfig(
    r=16,
    lora_alpha=32,                           # α = 2r
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, lora_cfg)
model.print_trainable_parameters()
# trainable params: 20,971,520 || all params: 7,262,703,616 || trainable%: 0.2887

20M trainable params on a 7B model. That's the whole trick.

4.2 Training

ds = load_dataset("tatsu-lab/alpaca", split="train").map(to_chatml)  # see S058
 
cfg = SFTConfig(
    output_dir="./mistral-lora-alpaca",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8,
    learning_rate=2e-4,                     # 40× higher than full SFT
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    bf16=True,
    logging_steps=25,
    save_steps=500,
    max_seq_length=2048,
    packing=True,
)
 
trainer = SFTTrainer(model=model, args=cfg, train_dataset=ds, tokenizer=tok)
trainer.train()
trainer.save_model("./mistral-lora-alpaca")   # saves only the LoRA weights (~80MB)

On a single A100 40GB, this runs in ~4 hours for 3 epochs. On an RTX 4090 24GB with QLoRA (S061), same time, same quality.

4.3 Loading a LoRA adapter for inference

from peft import PeftModel
 
base = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype="bfloat16")
model = PeftModel.from_pretrained(base, "./mistral-lora-alpaca")
model.eval()
 
# Generate as usual; at each forward pass, output = base(x) + (α/r) B A x

The base model is loaded once. The adapter is 80MB. You can load many adapters on the same base and switch between them at request time — this is how services like Predibase and OpenPipe run 1000+ fine-tunes on one GPU.


5 · Merging adapters for zero-latency inference

At inference time, LoRA adds a small overhead per forward pass (two extra matmuls per attention projection). For low-batch serving this is trivial; for high-throughput serving it costs 5–15% latency.

The merge trick: you can fold BA into W once, offline, producing a new W' = W + (α/r) · BA. Now inference has zero LoRA overhead — it's just a base model with slightly different weights.

merged = model.merge_and_unload()
merged.save_pretrained("./mistral-merged")

After merge, the LoRA structure is gone — you have a plain HuggingFace model with the same 7B parameters as the base, just tweaked. Any inference server (vLLM, TGI, llama.cpp) can serve it with no LoRA support required.

5.1 When NOT to merge

  • Multi-adapter serving. If you want to switch adapters per request, don't merge.
  • You'll further-tune. After merge, you can't easily continue LoRA on top; the boundary is gone.
  • Quantization mismatches. If your base is fp16 but you plan to serve int8, merging into fp16 first then quantizing may amplify rounding error. Better to merge into the target dtype directly.
War story War story

I once merged a LoRA adapter into a base model, quantized to int4 for llama.cpp, and shipped. On the test suite (5 examples) it worked. In production on real user traffic, it hallucinated much more than expected.

Root cause: the LoRA had learned some very small weight changes in a few attention heads (magnitude < 1e-3). Merged, these were tiny perturbations of the base. Quantized to int4 (which rounds to ~1/128 = 8e-3 precision), the entire LoRA signal was rounded away. I was effectively serving the raw base model.

Fix: keep the LoRA separate, load base in int4, apply LoRA in fp16 at forward time. peft supports this out of the box. Small latency hit, correct behavior.

Lesson: when quantizing after merge, verify that your LoRA weights survive the quantization. If |ΔW| < quant step, they don't.


6 · Diagram — the LoRA forward pass

Two paths. The wide path (W) is fixed at 16GB. The narrow path (BA) is trainable at ~4MB per layer. Backward pass only flows through the narrow path.


7 · LoRA in 2024–2025 — DoRA, rsLoRA, PiSSA, and the variants that actually helped

The 2021 LoRA paper was elegant but the community poked at it for four years. Here are the variants that beat vanilla LoRA on real benchmarks (as of mid-2025), and which one to reach for.

7.1 rsLoRA — rank-stabilized LoRA (Kalajdzievski, 2023)

Vanilla LoRA scales its update by α/r\alpha / r. Kalajdzievski showed this scaling is wrong at high rank — the update magnitude collapses as rr grows, so people who tried r=256 or r=512 were quietly training at a much smaller effective learning rate than they thought.

The fix is trivial: scale by α/r\alpha / \sqrt{r} instead. That's it. One character change (rmath.sqrt(r)) in peft's LoraConfig(use_rslora=True). On instruction-tuning benchmarks with r64r \geq 64, rsLoRA gains 1–3 points MT-Bench for free.

When to use: always, if r32r \geq 32. Zero downside.

Further reading: https://arxiv.org/abs/2312.03732

7.2 DoRA — weight-decomposed LoRA (Liu et al., NVIDIA 2024)

DoRA decomposes each pretrained weight matrix into a magnitude vector and a direction matrix, then applies LoRA only to the direction. The intuition: full fine-tuning tends to change both magnitude and direction, but LoRA (rank-limited) can only change direction meaningfully. By explicitly separating the two, DoRA gives the low-rank update the direction budget it needs and reserves a full-rank vector for magnitude changes.

Result: DoRA closes 60–80% of the remaining gap between LoRA and full fine-tuning across LLaMA, LLaVA, and VL-Bart benchmarks. Cost: ~15% more compute per step (one extra weight-normalization). Same memory footprint.

Usage: LoraConfig(use_dora=True) in peft >= 0.10. Costs one line.

When to use: you were about to reach for full fine-tuning because vanilla LoRA felt like it left performance on the table. Try DoRA first.

Further reading: https://arxiv.org/abs/2402.09353

7.3 PiSSA — principal singular values init (Meng et al., 2024)

Vanilla LoRA initializes AA with Kaiming-uniform and BB with zeros. Because B=0B=0, the initial adapter contributes nothing — you have to train BB from scratch. PiSSA initializes AA and BB from the top-rr singular vectors of the frozen pretrained weight W0W_0. So the adapter starts already "in" the model's principal directions.

Effect: much faster convergence (often 2–3× fewer steps to hit the same loss) and better final quality. Cost: one SVD per target layer at initialization, which takes seconds even for 8B models.

When to use: you have a tight compute budget and can't afford long SFT runs. PiSSA gets you 90% of the way in the first few hundred steps.

Further reading: https://arxiv.org/abs/2404.02948

7.4 LoRA+ — different LRs for A and B (Hayou et al., 2024)

A subtle finding: in vanilla LoRA, AA and BB should be trained with different learning rates. Specifically, ηB=16ηA\eta_B = 16 \cdot \eta_A works better across scales. Why? Because BB starts at 0 and AA starts near unit-scale — they need very different step sizes to move at comparable relative rates.

One-line change in peft. Gains 0.5–2 points across the LLaMA and Roberta suites. Trivial to try.

Further reading: https://arxiv.org/abs/2402.12354

7.5 The 2025 default recipe

Put it all together and the current "strong default" LoRA config for a 7B instruction fine-tune is:

from peft import LoraConfig, get_peft_model
 
cfg = LoraConfig(
    r=64,                       # up from 8/16 — memory is cheap now
    lora_alpha=128,             # = 2r
    lora_dropout=0.05,
    target_modules="all-linear",  # not just q/v — all linear layers
    use_rslora=True,            # rank-stabilized scaling
    use_dora=True,              # weight-decomposed (accepts 15% compute hit)
    init_lora_weights="pissa",  # SVD init
    bias="none",
)

Compared to the 2023 default (r=8, target=["q_proj","v_proj"]), this uses ~10× more trainable parameters (~40M for a 7B model, still tiny) but reliably matches or beats full fine-tuning on most SFT tasks.

7.6 War story — the target_modules trap

Early LoRA folklore (from the 2021 paper) was: only adapt WqW_q and WvW_v because that gave the best cost/quality tradeoff for GPT-3. This became gospel and every tutorial parroted it.

In 2023 a team I know spent two weeks trying to fine-tune Llama-2-7B for a legal-summarization task using target_modules=["q_proj", "v_proj"]. LoRA training "worked" (loss went down) but the model refused to shift its stylistic register from generic-web to legal-formal. They tried more epochs, more data, higher rank — nothing.

One of their engineers finally set target_modules="all-linear" (which in PEFT means q/k/v/o/gate/up/down). Same rank, same data, same time budget. Model started producing polished legal briefs after 300 steps.

The reason: the FFN layers hold most of the "style" knowledge in a transformer. Attention layers do information routing; FFNs do content transformation. If you want to shift style/domain, you MUST adapt the FFN. The 2021 paper's q/v-only recommendation was for factual task adaptation on GPT-3, not for style/persona shifts on modern LLMs.

Modern rule: target_modules="all-linear" for anything style/persona/domain related. Bump rr down if memory is tight; don't skip modules.


8 · Retention scaffold

Recall

Q1. What is the LoRA update rule in one equation?

Revealy = W x + (α/r) · B A x, where W is frozen (pretrained), A ∈ ℝr×d, B ∈ ℝd×r are trainable, B is zero-initialized, and A is Kaiming-initialized.

Q2. Why is B initialized to zero and not A?

RevealTo ensure the LoRA output is exactly zero at step 0, so the model behaves identically to the base at initialization. If B is zero, then BAx = 0 regardless of A. If we zeroed A instead, we would have the same effect — but the convention picks B because it typically has fewer rows than A has columns (numerical stability of first grad step).

Q3. Why does LoRA tolerate a 10–100× higher learning rate than full-parameter SFT?

RevealBecause only 0.1–1% of parameters are trainable, and they start from zero. The gradients flow into a much smaller subspace; a large LR that would nuke the full model just optimizes the small adapter fast. Also, the base weights (which have the model's competence baked in) never move.

Q4. Which modules should you target first when injecting LoRA into a transformer?

RevealAttention Q and V projections. This is the original LoRA paper's recommendation and remains the strong default. Adding K and O gives marginal gains at ~2× the params. Targeting MLPs doubles again with diminishing returns for SFT.

Q5. When should you NOT merge a LoRA adapter into the base weights?

Reveal(1) When you need to serve many adapters on one base model; (2) when you plan to continue-train the LoRA; (3) when you'll aggressively quantize (< int8) — the LoRA delta may be smaller than the quantization step and get rounded away. In case (3), keep LoRA separate and quantize only the base.

Stretch prompt

The original LoRA paper hypothesizes that ΔW is intrinsically low-rank. Design an experiment to measure the intrinsic rank of ΔW for a specific fine-tune. Hint: SVD the difference between fine-tuned and base weight matrices, look at the singular-value spectrum. What would you expect: a sharp knee at rank ~8, a gradual decay, or something else?

In your own words

Explain LoRA to a colleague who understands SFT but has never heard of adapters. In your explanation, address: (a) what problem it solves, (b) what the algorithm is, (c) what quality/cost tradeoff you should expect.

Spaced review

  • S015 — Linear layers. LoRA is a linear-algebraic refactoring; refresh matmul dimensions.
  • S058 — SFT loop. LoRA is a way to do SFT with less memory; the loss and training loop are unchanged.
  • S042 — KV cache. Post-LoRA, generation is still the same; only training changed.

Next session

S061 — QLoRA. LoRA got the trainable params down by 1000×. But the frozen base weights still need 16GB for a 7B model. QLoRA (Dettmers et al. 2023) quantizes the base to 4-bit (4× smaller), makes it 4GB, and combines with LoRA. Result: fine-tune a 7B on a consumer GPU. We'll cover NF4, double quantization, paged optimizers, and bitsandbytes.

Bring back tomorrow

  • The LoRA equation: y = Wx + (α/r) B A x, B zero-init.
  • α = 2r is a strong default.
  • LoRA LR is 10–100× higher than full-SFT LR (typically 2e-4).
Five things to remember about LoRA
    Session recall · click to reveal
    ★ = stretch question

    Previous: ← DL S059 · Next: DL S061 →