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.
🎯 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.
- 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:
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:
- B is initialized to zero.
- 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_outThat's the whole thing. 25 lines including whitespace.
2.1 Reading the forward pass
x @ self.lora_A.Tproduces shape(..., r)— projects to the low-rank space.@ self.lora_B.Tproduces shape(..., out_features)— projects back up.self.scaling * lora_outscales by α/r.- Summed with
base_out, which was computed by the frozenW 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 modelNow:
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.
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.288720M 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 xThe 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.
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 . Kalajdzievski showed this scaling is wrong at high rank — the update magnitude collapses as 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 instead. That's it. One character change (r → math.sqrt(r)) in peft's LoraConfig(use_rslora=True). On instruction-tuning benchmarks with , rsLoRA gains 1–3 points MT-Bench for free.
When to use: always, if . 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 with Kaiming-uniform and with zeros. Because , the initial adapter contributes nothing — you have to train from scratch. PiSSA initializes and from the top- singular vectors of the frozen pretrained weight . 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, and should be trained with different learning rates. Specifically, works better across scales. Why? Because starts at 0 and 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 and 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 down if memory is tight; don't skip modules.
8 · Retention scaffold
Recall
Q1. What is the LoRA update rule in one equation?
Reveal
y = 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?
Reveal
To 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?
Reveal
Because 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?
Reveal
Attention 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).