Fine-tuning and LoRA.

Fine-tuning is taking a model that already knows a lot and training it a little more so it gets good at one particular job, and the trick here is to leave the whole original model untouched and train a small, slim add-on next to it instead.

That trick decides whether you can do the job on your Spark at all. Retraining every number in an 8-billion-number model the ordinary way needs about as much memory as the whole machine has, before it even starts. The add-on version fits with room to spare, saves as a file small enough to email, and once you fold it back in, the model runs exactly as fast as before. That is how you would teach an 8B model to judge query-page relevance for your ranking work in a strict, machine-readable format, without it opening with “Sure! Here’s my assessment:” one time in twenty.

It builds on How much memory does a model need? and Why training needs so much memory, which give you the memory bill this trick dodges, and on Quantization for squeezing the frozen model into fewer bits. If you did Train your own draft model, this is a much cheaper way to run that same kind of retraining. Next, Why are evals separate from training? is about trusting the score you get, and the add-on comes back in DPO and GRPO, the cheap way to train models from preferences and rewards on one machine.

Today

• Work out, in bytes, why retraining every number of an 8B model does not fit on the Spark, and why the slim add-on version (LoRA), and the version that also squeezes the frozen model into 4 bits (QLoRA), do.
• Draw the LoRA picture from memory: the big frozen d × d grid of numbers, plus a thin B times a flat A beside it, the count of numbers that buys you (2dr against d²), and why B starts at all zeros.
• Write that add-on layer from scratch in PyTorch, plug it into a real model, and prove that only A and B learn.
• Fine-tune a 1-8B model on the Spark with LoRA and QLoRA for a ranking-shaped task, test it properly against the untouched model, fold the add-on back in, and serve it from your own server.
• Say clearly what this trick is good at (a format, a style, one narrow task) and what it is bad at (absorbing lots of new knowledge), with the paper that measured it.

You will want How much memory does a model need? for bytes per number, and ideally Why training needs so much memory for the 16-bytes-per-number bill. Quantization helps for the 4-bit part. Plan on one evening for the from-scratch layer on your Mac, then one or two sessions on the Spark.

The Problem

Your cross-encoder ranks documents well, but you want something it can’t do: a model that reads a query and a page and returns a strict, machine-readable judgement, {"label": 2, "reason": "..."}, for labelling training data or for a second-stage check. You try prompting Llama 3.1 8B Instruct. It is right most of the time. It also, about one time in twenty, opens with “Sure! Here’s my assessment:”, forgets a field, or invents a label 4 on a 0-3 scale. Your parser falls over, at 3 a.m., on the one batch that mattered.

The obvious fix is to train the model on a few thousand examples of exactly the output you want. That’s fine-tuningfine-tuningTaking a model that is already trained (pretrained) and training it a bit more on a smaller, specific dataset so it gets good at one job.. So you do the arithmetic from the training-memory lesson. Training with Adam in mixed precision costs about 16 bytes per parameter before you store a single activation. 8 billion parameters × 16 bytes = 128 GB. The Spark has 128 GB of memory in total (spec-sheet number), shared with the operating system and the CPU. It doesn’t fit. Not “fits badly”: doesn’t fit.

And suppose it did. Every task you fine-tune for gives you a new 16 GB copy of the model. A relevance judge, an extractor, a query rewriter: three copies, 48 GB of disk, and if you want to serve all three you need three models in memory.

Here’s the strange thing. The change fine-tuning makes to the weights turns out to be tiny and simple. You’re paying to update 8 billion numbers when the change you need could often be written down with a few million.

The Idea

Fine-tuning is a small edit to a big matrix

Strip a transformer down and most of its parameters live in a handful of big linear layerlinear layerA layer that computes y = Wx (plus an optional bias): multiply the input vector by a weight matrix. In PyTorch, nn.Linear.s per block: the attention projections q_proj, k_proj, v_proj, o_proj, and the MLP’s gate_proj, up_proj, down_proj. Each is a matrix W, and the layer computes h = W x.

Fine-tuning starts from the pretrained W and nudges it. After training you have a new matrix, which you can always write as the old one plus a difference:

W_new = W + ΔW

Full fine-tuning learns ΔW with no restrictions: every one of its entries is free. For a 4096 × 4096 projection that’s 16.8 million numbers, for one matrix, in one layer. LoRA’s whole idea is to put a restriction on the shape of ΔW so it can be stored and trained with far fewer numbers.

The bill for doing it the full way

Before the trick, the cost of not using it. When you train every weight with AdamAdamThe standard optimizer for transformers. For every parameter it keeps two running averages: of the gradient (momentum) and of the squared gradient (to scale each step). AdamW is Adam with decoupled weight decay. in mixed precisionmixed precisionRunning the forward and backward passes in a 16-bit format (bf16) for speed, while keeping a 32-bit “master” copy of the weights so tiny updates don’t round away., every parameter carries:

Implementations differ a little (some keep fp32 gradients, some drop the bf16 copy), so treat 16 as the standard round number; Why training needs so much memory takes it apart. On top of this you need activationsactivationsThe intermediate outputs of every layer from the forward pass, kept in memory because the backward pass needs them to compute gradients. Grows with batch size × sequence length., which grow with batch size times sequence length.

Now look at what a frozen weight costs. If a parameter doesn’t train, it needs no gradient, no master copy and no optimizer state. It costs just its 2 bf16 bytes, or about half a byte if you store it in 4 bits. So the game is: freeze almost everything, and make the part that trains tiny.

The “~5.7 GB” for QLoRA is my estimate, not a measurement: the transformer layers at a little over 4 bits per weight (~3.6 GB), plus the input embedding and output head, which are normally left in bf16 (~1 GB each for Llama 3’s 128k vocabulary). You’ll measure the real number.

The low-rank bet

Here is the restriction. Instead of learning ΔW (shape d × d) directly, learn two thin matrices and multiply them:

ΔW = B · A        B is d × r,  A is r × d,  and r is small (8, 16, 64)

Think of what B · A does to an input x. First A squeezes the d-dimensional x down to just r numbers. Then B spreads those r numbers back out to d dimensions. Whatever the update does, it has to pass through an r-number bottleneck. So the update can only push the output in r independent directions. That is exactly what the rankrankThe number of independent directions a matrix can output. A d × d matrix has rank at most d. A product B·A with an r-wide middle has rank at most r, however big d is. of a matrix means, and why B · A is called a low-ranklow-rankA matrix whose rank is much smaller than its size: it can be written exactly as a product of a tall-thin and a short-wide matrix. update.

Count the numbers. B has d × r, A has r × d, so together 2dr. Against d² for the full matrix:

For a non-square matrix (input size k, output size d, like Llama’s MLP at 4096 → 14336) the same logic gives r × (d + k). Summed over every linear layer of Llama 3.1 8B at r = 16, that’s 41,943,040 trainable parameters, about 0.52% of the model. Only q_proj and v_proj at r = 16: 6.8 million, under 0.1%.

The bet is that the update fine-tuning needs is close to low-rank, so this restriction costs little. Whether that’s true depends on the update. Play with it.

Three things to notice while you drag the slider.

1. The shapes. At r = 4, B and A are slivers next to W, and their combined count is a third of d² even at this toy size. At d = 4096 the same r is 0.2%. The savings grow with d, because d² grows faster than 2dr.
2. The error curve’s shape is the whole story. If the update has a few strong directions, the error falls off a cliff at the right rank and then sits at zero. Rank beyond that is wasted. If the update has energy spread over many directions, the curve is a slow slope, and any small rank leaves most of it on the table.
3. The spectrum bars. The singular valuessingular valuesFor any matrix, the strengths of its independent directions, sorted largest first (from the SVD). Rank = how many are non-zero. Keeping the top r gives the best rank-r approximation (the Eckart-Young theorem). are those strengths, sorted. LoRA can only ever capture the top r of them (the blue ones). Everything grey is gone.

Why a tiny rank is often enough

The deeper reason comes from a 2020 paper by Armen Aghajanyan, Luke Zettlemoyer and Sonal Gupta, “Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning”. They asked: how many free numbers do you actually need to fine-tune a pretrained model to good accuracy?

Their method: freeze the model, pick a random way to map a small vector of k numbers to a change in all the weights (a fixed random projection), and train only those k numbers. Then find the smallest k that still gets you 90% of full fine-tuning’s accuracy. That k is the task’s intrinsic dimensionintrinsic dimensionThe smallest number of free parameters that, mapped into the full weight space, can still solve the task to a set fraction (here 90%) of full fine-tuning’s quality. Li et al. 2018 introduced the measurement..

The numbers were startling. For a RoBERTa model on a paraphrase task (MRPC), about 200 trainable numbers, randomly projected into all of its hundreds of millions of weights, got 90% of full fine-tuning’s performance. And bigger, better-pretrained models had smaller intrinsic dimension: pretraining makes the model so ready for language tasks that adapting it is a short walk in a few directions.

LoRA turns that into an engineering choice: instead of a random projection, a structured one (low rank, per matrix) that is cheap to train and can be multiplied out at the end.

The LoRA layer, piece by piece

Here’s what one LoRA-wrapped linear layer computes:

h = W x  +  (α / r) · B (A x)

Three design details hide in that line, and each one is there for a reason.

B starts at zero. With B = 0, the side path outputs zero, so h = W x exactly: on step 0 the model is the base model, with all its abilities. Training starts from a known-good point and moves away only as far as the data pushes it. A is initialised randomly (the paper used a Gaussian; Hugging Face’s PEFT library uses the same Kaiming-uniform init as nn.Linear). Why not both zero? Look at the gradients. The gradient for B is proportional to A x, and the gradient for A is proportional to Bᵀ times the incoming gradient. If both start at zero, both gradients are zero, forever, and nothing ever trains. One of them must be random so the other gets a signal. You’ll see this exact behaviour in Break It: at step 0, A’s gradient is exactly zero and B’s is not.

The scale α / r. α (alpha) is a fixed number you pick, usually r or 2r. The point of dividing by r is to keep the size of the update roughly steady as you change r, so that you can try rank 8 and rank 32 without retuning the learning rate from scratch. The LoRA paper simply set α to the first r they tried and never tuned it. A later analysis, rsLoRArsLoRARank-stabilised LoRA (Kalajdzievski, 2023): scale by α/√r instead of α/r, which keeps learning stable at large ranks. A flag in PEFT (use_rslora)., argued that α / √r behaves better at large ranks, and PEFT supports both. In practice: fix α relative to r (say α = 2r), and tune the learning rate.

LoRA dropoutLoRA dropoutDropout applied to the input of the A·B side path only (not to the frozen path), to regularise the adapter. Typical values 0 to 0.1.. Optional dropout on the input to the side path, usually 0.05 to 0.1, as a mild guard against overfitting small datasets.

Which matrices to adapt

You choose which linear layers get a LoRA wrapper. The LoRA paper, short on memory for GPT-3, adapted only the query and value projections (W_q, W_v), and found that spreading a fixed budget over both beat putting it all in one. That became the default in a lot of tutorials.

It’s not the best default anymore. QLoRA found that adapting every linear layer, attention and MLP, was needed to match full fine-tuning quality, and that once you do, the rank matters much less. The “LoRA Learns Less and Forgets Less” study found the same, and Thinking Machines’ 2025 write-up “LoRA Without Regret” found that MLP layers matter most: attention-only LoRA underperformed even at matched parameter count. Most of a transformer’s parameters live in the MLP, so that’s not surprising.

Practical default: target all linear layers in the blocks (PEFTPEFTHugging Face’s library for parameter-efficient fine-tuning: LoRA and its relatives, wrapped around Transformers models. accepts target_modules="all-linear"), leave the embeddings and the output head frozen, and start at r = 16.

Merging: an adapter that costs nothing to run

After training you have W (frozen, unchanged) and a small B, A per adapted layer. You can serve that as is, computing the side path on every token. Or you can multiply it out once:

W_merged = W + (α / r) · B · A

W_merged has exactly the shape of W. The model is now an ordinary model, same architecture, same speed, nothing extra on the forward path. That’s the property that made LoRA beat adapters: zero added latency at inference, because there’s nothing left to add. PEFT calls it merge_and_unload(). You can also go back: subtract the same product and you have the base model again (up to rounding).

Merging has one trap with QLoRA. You trained the adapter against a 4-bit copy of the base. If you merge into the 4-bit weights, the sum has to be re-quantized and much of your tiny update rounds away. The usual practice is to load the base in bf16, merge the adapter into that, and save a bf16 model. That model is very slightly different from the one you trained (bf16 base instead of 4-bit), so always run your eval on the merged model, not just the training checkpoint.

Serving many adapters on one base

Merging bakes one task in. But the other choice, not merging, has a superpower: the base weights are shared, and each adapter is tiny (tens of megabytes). So one server can hold one copy of the base and hundreds of adapters, and pick an adapter per request.

The hard part is the per-row step: a batch where row 1 uses adapter “judge” and row 2 uses “extract” can’t be one plain matmul. PunicaPunicaA 2023 system (Chen et al.) for serving many LoRA adapters on one GPU. Introduced a kernel (SGMV) that applies a different adapter to each request in a batch. wrote a custom kernel that does a different small matmul per segment of the batch, and S-LoRAS-LoRA“Serving Thousands of Concurrent LoRA Adapters” (Sheng et al., 2023). Keeps adapters in CPU memory, pages active ones onto the GPU in a pool shared with the KV cache, and batches requests across adapters. added paging: adapters live in CPU memory and move into a GPU pool shared with the KV cache as requests arrive, the same paging idea as vLLM’s KV blocks from What’s inside a serving engine?. vLLM supports this directly: enable LoRA, register adapters by name, and a request picks one through its model field (flags like --enable-lora and --lora-modules name=path; check vLLM’s LoRA docs for your version). The side path now runs on every token, so unmerged adapters cost a little speed: the price of flexibility.

QLoRA: freeze the base in 4 bits

LoRA got rid of the gradients and optimizer states for frozen weights. The frozen weights themselves are still 2 bytes each: 16 GB for 8B, 130 GB for 65B. In 2023 Tim Dettmers, Artidoro Pagnoni, Ari Holtzman and Luke Zettlemoyer asked: the base is frozen, so why store it in 16 bits at all? QLoRAQLoRAQuantized LoRA (Dettmers et al., 2023): store the frozen base in 4-bit NF4, dequantize each weight to bf16 on the fly for the matmul, and train bf16 LoRA adapters on top. Fine-tuned a 65B model on a single 48 GB GPU. fine-tuned a 65B model on a single 48 GB GPU, and matched 16-bit fine-tuning quality on their benchmarks. Three pieces made it work.

NF4. From Quantization you know a 4-bit format has 16 levels and the question is where to put them. Pretrained weights are roughly bell-shaped (normal) around zero. NF4NF44-bit NormalFloat. The 16 levels are placed at quantiles of a normal distribution, so each level is used about equally often by normally distributed weights. Each block of 64 weights has its own scale (its absolute max). places the 16 levels at the quantiles of a normal distribution, so each level gets used by roughly the same number of weights: dense near zero where most weights live, sparse in the tails. Each block of 64 weights gets its own scale. For the matmul, the block is unpacked to bf16 on the fly; the adapter math is in bf16 throughout. Gradients flow through the 4-bit weights (to reach earlier layers) but never update them.

Double quantization. Every block of 64 weights stores one fp32 scale: 32 bits / 64 weights = 0.5 extra bits per weight. That’s 4 GB for a 65B model, just for scales. QLoRA quantizes the scales too: to 8 bits, in blocks of 256 scales, each with its own fp32 second-level scale. Cost per weight becomes 8/64 + 32/(64 × 256) ≈ 0.127 bits, saving about 0.37 bits per weight. About 3 GB on a 65B model.

Paged optimizers. Memory use during training spikes, for example on an unusually long batch, and a spike at the wrong moment is an out-of-memory crash. QLoRA put the optimizer states in NVIDIA unified memoryunified memoryCUDA memory that can automatically move pages between GPU and CPU memory as they’re touched, like an OS swapping pages to disk., so on a spike they get paged out to CPU RAM and paged back when needed, like the OS swapping from How does memory really work?. On the Spark, CPU and GPU already share one pool of memory, so there’s less to page between, but the same option protects you from a crash.

The cost: every frozen weight is dequantizeddequantizedUnpacked from its stored low-bit form back to a normal float (here bf16), using its block’s scale, just before the matmul that needs it. on every forward and backward pass, so QLoRA steps are slower than bf16 LoRA. On the Spark, bf16 LoRA on an 8B fits easily; QLoRA is a choice until you try a 70B.

What LoRA can and can’t learn

Here’s the honest part, and it decides when to reach for LoRA at all.

In 2024, Dan Biderman, Jacob Portes, Jose Javier Gonzalez Ortiz and colleagues (Columbia and Databricks) ran the careful comparison in “LoRA Learns Less and Forgets Less”. They trained Llama-2-7B on code and on math, in two regimes: continued pretrainingcontinued pretrainingTraining further on large amounts of raw text from a new domain (billions of tokens), to teach the model knowledge, not just a format. (billions of unlabeled tokens) and instruction fine-tuninginstruction fine-tuningTraining on prompt-response pairs so the model learns to follow instructions or a task format. Usually thousands to a few hundred thousand examples. (prompt-response pairs). They compared full fine-tuning with LoRA at various ranks. What they found:

• LoRA learns less. On the target domain, LoRA substantially underperformed full fine-tuning, most of all in continued pretraining on code. It was much closer in instruction fine-tuning. The full fine-tuning updates had high rank, 10-100× more than typical LoRA settings. A low-rank adapter can’t hold a domain’s worth of new facts.
• LoRA forgets less. Measured on what the base model could already do (language understanding and common-sense benchmarks such as HellaSwag, ARC and WinoGrande), LoRA kept more of it than full fine-tuning did. A small edit breaks less.
• Their advice: target all modules, use a higher learning rate than full fine-tuning, and use LoRA where the goal is adapting behaviour, not absorbing a corpus.

That matches Figure 002. A new output format, a style, a narrow classification or extraction task: “few directions”, and LoRA is close to full fine-tuning. Everything about your company’s internal APIs, or a new language: “many directions”, and LoRA leaves much of it out. For knowledge, retrieval (facts in the prompt) usually beats fine-tuning anyway. Thinking Machines’ “LoRA Without Regret” (2025) refines the boundary: with all layers adapted and the right learning rate, LoRA matched full fine-tuning on small-to-medium post-training datasets, and fell behind once the dataset carried more information than the adapter had capacity for.

The flip side, forgetting, has a name. catastrophic forgettingcatastrophic forgettingWhen training on a new task overwrites what a network knew before, so its ability on old tasks drops sharply. Named by McCloskey and Cohen (1989) for small networks; still very real for LLMs. is what happens when your fine-tune gets great at relevance labels and quietly gets worse at everything else: it stops following other instructions, its general knowledge slips, it answers every question with a relevance label. Lower rank, fewer steps, a lower learning rate, and mixing a little general instruction data into your training set all reduce it. And you only know how much you lost if you measure it, which is why the eval below has two halves.

Data: chat templates and loss masking

The model learns from tokens, so the most common fine-tuning bugs are about tokens, not math.

Use the model’s own chat templatechat templateThe exact text layout, including special marker tokens, that an instruct model was trained to see conversations in. Stored with the tokenizer and applied by tokenizer.apply_chat_template.. An instruct model was trained to see conversations wrapped in specific special tokens. For Llama 3 a user turn looks like <|start_header_id|>user<|end_header_id|>, a blank line, the text, then <|eot_id|>. Other model families use other markers. If you train on raw “Query: … Answer: …” text and then serve through the chat endpoint (which applies the template), the model sees a format at inference that it never saw in training. Always build training examples with tokenizer.apply_chat_template(...), and serve the same way.

Mask the prompt out of the loss. In a training example, the prompt (system message, query, a long document) might be 500 tokens and the answer 15. If you train on every token, 97% of the gradient goes into teaching the model to predict your prompts, which it doesn’t need to generate. You want loss only on the answer. In Hugging Face models you do that with the labels: copy the input token IDs into labels, then set every prompt position to -100. The cross-entropy loss ignores any position labelled -100 (ignore_index in PyTorch’s F.cross_entropy). Keep the end-of-turn token in the loss, or the model never learns to stop.

One subtle trap when you build the mask: tokenize the prompt and the full conversation separately and check that the full one starts with the prompt’s tokens. Tokenizers can merge characters across the boundary, so “prompt length” in tokens isn’t always where you think it is. The from-scratch code below asserts this.

Libraries do the masking for you (TRL’s SFTTrainer has options for training only on the completion or the assistant turns; their names have changed across versions, so check its docs). Do it by hand once anyway, so you can check the library is doing it.

Hyperparameters, and how overfitting shows up

The knobs, with starting points that are widely used, not laws:

overfittingoverfittingWhen a model fits the training examples better and better while getting worse on new ones: it memorises rather than learns the pattern. in fine-tuning looks like this:

• Train loss keeps dropping, held-out loss turns up. The classic. Evaluate on a held-out set every few dozen steps and keep the checkpoint where held-out loss (or better, your task metric) was best.
• Train loss near zero. On a few thousand examples, a loss of 0.01 means the model has memorised them. That’s almost never what you want.
• Outputs copy training examples. Reasons that repeat word-for-word from training data, or the same phrase on every input.
• Everything becomes your task. Ask it “what’s the capital of France?” and it returns {"label": 0, ...}. That’s forgetting showing up as overfitting to the format.

Evaluating a fine-tune

“The loss went down” is not an evaluation. A proper one has two halves.

1. Did it learn the task? A held-out setheld-out setExamples kept completely out of training (and out of any decisions you tune on), used only to measure. If you tune on it, it stops being held out. scored with the metric you actually care about: label accuracy and macro-F1macro-F1F1 score (the balance of precision and recall) computed per class and then averaged, so rare classes count as much as common ones. Plain accuracy can hide a model that always predicts the majority label. for relevance grades, exact-match and field-level accuracy for extraction, plus “fraction of outputs that parse”. Compare against the base model with your best prompt: if prompting already gets 95%, a fine-tune that gets 96% may not be worth owning. Split carefully. For ranking data, split by query, not by row: if the same query appears in train and test with different documents, the model can learn the query, and your test score lies. (Why are evals separate from training? goes deeper.)

2. Did it break anything else? A general-capability check, run identically on base and fine-tune: a small slice of standard benchmarks (EleutherAI’s lm-evaluation-harness runs MMLU, ARC, HellaSwag and friends) and a fixed set of 30-50 everyday prompts you read side by side. You’re looking for drops, not gains.

And check it on the model you’ll actually serve: the merged one, through your server, with the same chat template and decoding settings.

Build It

Two halves. On your Mac: build LoRA yourself, small enough to understand every number. On the Spark: use the real library on a real model for a task from your work, then merge, serve and measure.

lora/
  lora_layer.py        # your LoRALinear + inject() + merge()
  toy_train.py         # inject into SmolLM2-135M, train a toy extraction task
  data/build.py        # your task data -> train/val/test JSONL (split by query)
  data/*.jsonl
  peft_train.py        # PEFT LoRA / QLoRA / full FT, one script, flags
  eval.py              # task eval + general check, any model dir or adapter
  merge.py
  results/*.json

cd ~/workbench && mkdir -p lora && cd lora
python3 -m venv .venv && source .venv/bin/activate
pip install torch transformers
python - <<'EOF'
from transformers import AutoModelForCausalLM, AutoTokenizer
mid = "HuggingFaceTB/SmolLM2-135M-Instruct"
tok = AutoTokenizer.from_pretrained(mid)
m = AutoModelForCausalLM.from_pretrained(mid)
print(sum(p.numel() for p in m.parameters()))   # ~134.5M
print(m.model.layers[0])                         # find the nn.Linear names
EOF

# lora_layer.py
import math
import torch
import torch.nn as nn

class LoRALinear(nn.Module):
    def __init__(self, base: nn.Linear, r: int = 8, alpha: float = 16, dropout: float = 0.0):
        super().__init__()
        # keep `base`, freeze its weight (and bias, if any)
        # A: (r, in_features), random init: nn.init.kaiming_uniform_(A, a=math.sqrt(5))
        # B: (out_features, r), zeros
        # self.scale = alpha / r
        ...

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # base(x) + scale * (dropout(x) @ A^T @ B^T)
        # order matters for cost: (x @ A^T) first makes an r-wide intermediate.
        ...

    @torch.no_grad()
    def merged(self) -> nn.Linear:
        """Return a plain nn.Linear whose weight is W + scale * B @ A."""
        ...

def inject(model: nn.Module, targets: tuple[str, ...], r: int, alpha: float) -> int:
    """Replace every nn.Linear whose attribute name is in `targets` with LoRALinear.
    Freeze every other parameter in the model. Return the expected count of trainable
    params from the formula r * (in_features + out_features), summed."""
    ...

def merge_all(model: nn.Module) -> None:
    """Replace every LoRALinear in the model with its merged nn.Linear."""
    ...

# toy_train.py (the data part; take it as is)
import json, random
names = ["Ana", "Bo", "Chen", "Dara", "Eli", "Femi", "Gus", "Hana", "Ivo", "Jun", "Kai", "Lena"]
cities = ["Lisbon", "Osaka", "Lima", "Oslo", "Accra", "Pune", "Quito", "Riga", "Seoul", "Tunis"]
templates = ["{n} is {a} and lives in {c}.", "Meet {n}, a {a}-year-old from {c}.",
             "{c} resident {n} turned {a} this year.", "At {a}, {n} moved to {c}."]

def example(rng: random.Random) -> tuple[str, str]:
    n, c, a = rng.choice(names), rng.choice(cities), rng.randint(18, 90)
    prompt = "Extract name, age, city as JSON: " + rng.choice(templates).format(n=n, a=a, c=c)
    answer = json.dumps({"name": n, "age": a, "city": c})
    return prompt, answer

def encode(tok, prompt: str, answer: str) -> tuple[list[int], list[int]]:
    """input_ids for the full chat (user prompt + assistant answer + end token),
    and labels with every prompt position set to -100.
    Build the prompt with tok.apply_chat_template(..., add_generation_prompt=True, tokenize=False)
    and assert the full sequence starts with the prompt's token ids."""
    ...

def collate(batch, pad_id: int):
    """Right-pad input_ids with pad_id, labels with -100, and build attention_mask."""
    ...

def exact_match(model, tok, examples, device) -> float:
    """Greedy-generate up to 40 new tokens for each prompt; fraction whose
    stripped output equals the answer string exactly."""
    ...

pip install transformers peft accelerate datasets
pip install bitsandbytes   # for QLoRA; see below
python -c "import peft, transformers; print(peft.__version__, transformers.__version__)"
python -c "import bitsandbytes as bnb, torch; print(bnb.__version__); \
  l = bnb.nn.Linear4bit(64, 64, quant_type='nf4', compute_dtype=torch.bfloat16).cuda(); \
  print(l(torch.randn(2, 64, device='cuda', dtype=torch.bfloat16)).shape)"

import torch
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

lora_cfg = LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05,
                      target_modules="all-linear", task_type="CAUSAL_LM")

if mode == "qlora":
    bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
                             bnb_4bit_use_double_quant=True,
                             bnb_4bit_compute_dtype=torch.bfloat16)
    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, quantization_config=bnb, device_map="cuda")
    model = prepare_model_for_kbit_training(model)   # also enables gradient checkpointing
else:
    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.bfloat16).to("cuda")

if mode in ("lora", "qlora"):
    model = get_peft_model(model, lora_cfg)
    model.print_trainable_parameters()   # check against your formula

Measure

Fill in. Predictions first: memory from bytes per parameter (16 for trained, 2 for frozen bf16, ~0.5 for frozen 4-bit, plus activations, which you estimate or measure with a tiny batch), step time from about 6 FLOPs per parameter per token for full FT and about 4 for LoRA, divided by the bf16 matmul throughput you actually measure on the Spark (time one big torch.matmul in bf16 first).

For reference, the toy on my M-series MacBook (SmolLM2-135M-Instruct in fp32 on MPS, all seven projections at r = 8, batch 16, 250 steps, lr = 2e-4, α = 2r). Machine was busy with other work, so treat the times as rough:

Both trainable counts matched the formula exactly. The base model’s zero is instructive: in the outputs I read, it got the name, age and city right but pretty-printed the JSON across several lines, so exact match failed. Its starting loss of 0.24 says the same thing, it already “knew” the answer. What LoRA taught it was the format. That’s the “few directions” case, which is why even 0.34% of the parameters was enough here. Notice also that the q,v run was not 5× faster for 5× fewer trainable parameters: the frozen layers’ forward and backward passes dominate the step.

What the numbers should tell you:

Trainable parameters match your formula exactly. If PEFT’s printout disagrees with your hand count, you targeted something you didn’t mean to (or missed something). This is the cheapest bug check in the whole lesson.

Memory: LoRA’s peak is mostly the frozen base plus activations. For the 8B in bf16, expect ~16 GB of weights plus activations that depend heavily on batch size × sequence length and whether gradient checkpointinggradient checkpointingSaving only some activations during the forward pass and recomputing the rest during backward. Trades roughly one extra forward pass of compute for much less activation memory. is on. The adapter’s own weights, gradients and Adam states are under 1 GB. QLoRA should cut the weights part to roughly a third. Full FT on the 1B lands near 20 GB plus activations. If your measured number is far above prediction, the gap is almost always activations: try halving the sequence length and watch it move.

Step time: LoRA is not 200× faster. Expect LoRA to take roughly two thirds of a full fine-tuning step at the same model size, and QLoRA to be slower than bf16 LoRA because of dequantizing every weight twice per step. If QLoRA is much slower, check that the compute dtype is bf16 and not fp32.

Task metric: LoRA ≈ full FT, both ≫ base on parse rate. For a narrow format-and-label task, expect LoRA and full fine-tuning on the 1B to land close to each other, and the fine-tunes to fix the format failures almost completely. Label accuracy improves less dramatically than parse rate: the format is “few directions”, judgement is harder.

General check: full FT forgets more. If anything drops, it should drop more for full FT than LoRA, and more with more epochs. Look at your 30 everyday prompts: does the fine-tune answer them as relevance judgements?

Merged vs base through the server: identical speed. If TTFT or tokens/s differs measurably, something isn’t merged (you’re still running the PEFT wrapper) or you saved in a different dtype.

Break It

1. Initialise both A and B to zero. In your LoRALinear, zero A too, and train the toy. The loss never moves, not even a little. Print the gradients: both are exactly zero at every step, because each one’s gradient is multiplied by the other. Then do the opposite: initialise B randomly too. Training runs, but check 3 from t2 fails: at step 0 the model is no longer the base model, it’s the base model plus a random perturbation in every layer, and early training spends its effort undoing damage. Zero B is the cheap way to start exactly where you want.

2. Use the wrong learning rate in both directions. Train the toy LoRA at lr = 2e-5 (a typical full fine-tuning rate): after 250 steps it has barely learned the format. Then full fine-tune SmolLM2 at lr = 2e-4 (a typical LoRA rate): loss spikes or goes to nan, or it learns the task and forgets how to say anything else. The same number means different things depending on how many parameters it moves and where they start.

3. Forget the loss mask. Set labels = input_ids for the whole sequence in the toy (or on your real task). The loss you print is now mostly about predicting your prompt template, which falls fast and looks great. The task metric improves much more slowly per step. Then the inverse bug: mask the end-of-turn token too. The model gets the JSON right and keeps going, generating a second turn or rambling until max_new_tokens.

4. Train in one format, serve in another. Train the toy on raw text ("Extract ...: <sentence>\n<answer>", no chat template), then evaluate through apply_chat_template (which is what your server does). Exact match drops sharply. Nothing is wrong with the weights; the input at serving time just isn’t the input the adapter learned to respond to.

5. Merge into the 4-bit weights. Take your QLoRA adapter, merge it into the 4-bit model and re-quantize (or simulate: merge in bf16, then round-trip the merged weights through NF4), and run the task eval. Compare with merging into a bf16 base. The update is small next to the 4-bit grid’s step size, so a lot of it rounds away, and the metric falls back toward the base model.

Explain It Back

Answer out loud, in plain words, as if in an interview. Then check.

Go Deeper

• Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, LoRA: Low-Rank Adaptation of Large Language Models (2021). The original. Read section 4 (the method, init and scaling) and section 7, where they measure the rank of the learned updates and what directions they amplify.
• Armen Aghajanyan, Luke Zettlemoyer, Sonal Gupta, Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning (2020). The “why it works” paper: a couple of hundred well-chosen numbers can fine-tune a big model, and bigger models need fewer.
• Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, Luke Zettlemoyer, QLoRA: Efficient Finetuning of Quantized LLMs (2023). NF4, double quantization and paged optimizers, plus the finding that adapting all linear layers matters more than rank.
• Dan Biderman, Jacob Portes, Jose Javier Gonzalez Ortiz, et al., LoRA Learns Less and Forgets Less (2024). The careful, honest comparison with full fine-tuning on code and math. The spectrum analysis is the best picture of when low rank is not enough.
• John Schulman and Thinking Machines Lab, “LoRA Without Regret” (2025, on the Thinking Machines blog). Practical, experiment-driven guidance: all layers, ~10× learning rate, when LoRA matches full fine-tuning and when capacity runs out.
• Ying Sheng et al., S-LoRA: Serving Thousands of Concurrent LoRA Adapters (2023). How to batch requests across many adapters with paging; read alongside the vLLM docs page on LoRA adapters.
• Hugging Face PEFT documentation, LoRA section (huggingface.co/docs/peft). The real API for LoraConfig, merge_and_unload, and QLoRA via bitsandbytes; check it against the code shapes above.

Connections

This builds on How much memory does a model need? (bytes per parameter for frozen weights) and Why training needs so much memory (the 16-byte bill LoRA dodges), and borrows NF4’s idea from Quantization: put the levels where the weights are. The multi-adapter serving picture is the paging idea from What’s inside a serving engine? applied to adapters instead of KV blocks, and it’s why prefix caching has to include the adapter ID in its cache key.

Next, the evaluation half deserves its own lessons: Why are evals separate from training? and Statistics for evals (is 94% vs 92% on 500 examples even a real difference?). And LoRA is the default way to run preference and RL fine-tuning on one machine, so it comes back in DPO and GRPO, where “forgets less” turns out to be a feature.