Fine-Tuning with PEFT — LoRA and Adapters

A senior engineer at Stripe knows everything about payments. You want to teach them your company's specific internal processes. You do not re-hire them and retrain them from scratch — you give them a small notebook of company-specific notes to carry alongside their existing expertise. LoRA is that notebook — small, lightweight, task-specific, sits alongside the frozen base model.

At inference time: base model knowledge + LoRA notebook = specialised expert. You can swap notebooks — same base model, different LoRA adapters for different tasks. One GPU, many specialists.

import torch
import torch.nn as nn
import numpy as np

# ── LoRA layer from scratch — see exactly what happens ────────────────
class LoRALinear(nn.Module):
    """
    Drop-in replacement for nn.Linear that adds LoRA adaptation.
    The original weight W is frozen. Only A and B are trained.
    """
    def __init__(self, linear: nn.Linear, rank: int = 8, alpha: float = 8.0):
        super().__init__()
        self.linear = linear
        self.rank   = rank
        self.alpha  = alpha
        self.scale  = alpha / rank

        d_out, d_in = linear.weight.shape

        # Freeze the original weight
        for param in self.linear.parameters():
            param.requires_grad = False

        # LoRA matrices — small and trainable
        self.lora_A = nn.Parameter(torch.randn(rank, d_in) * 0.02)
        self.lora_B = nn.Parameter(torch.zeros(d_out, rank))   # zeros → no change at start

    def forward(self, x):
        # Original frozen output
        base_out  = self.linear(x)
        # LoRA update: x @ A.T @ B.T * scale
        lora_out  = (x @ self.lora_A.T @ self.lora_B.T) * self.scale
        return base_out + lora_out

# ── Demonstrate parameter counts ──────────────────────────────────────
torch.manual_seed(42)
d_model = 768
original_linear = nn.Linear(d_model, d_model)

for rank in [4, 8, 16, 32]:
    lora_linear = LoRALinear(original_linear, rank=rank)

    total    = sum(p.numel() for p in lora_linear.parameters())
    trainable = sum(p.numel() for p in lora_linear.parameters() if p.requires_grad)
    frozen   = total - trainable
    pct      = trainable / total * 100

    print(f"rank={rank:2d}: total={total:,}  frozen={frozen:,}  "
          f"trainable={trainable:,}  ({pct:.2f}%)")

# ── Verify: LoRA output equals base at initialisation ─────────────────
lora = LoRALinear(original_linear, rank=8)
x    = torch.randn(4, d_model)

base_out = original_linear(x)
lora_out = lora(x)

diff = (base_out - lora_out).abs().max().item()
print(f"
Max diff between base and LoRA at init: {diff:.2e}  ← should be ~0")
print("B is initialised to zeros → LoRA adds nothing at the start")
print("Training gradually moves A and B to add task-specific adjustments")

# pip install peft transformers accelerate

from peft import (
    LoraConfig, get_peft_model, TaskType,
    PeftModel, prepare_model_for_kbit_training,
)
from transformers import (
    AutoTokenizer, AutoModelForSequenceClassification,
    AutoModelForCausalLM, TrainingArguments, Trainer,
    DataCollatorWithPadding,
)
from datasets import Dataset
import evaluate
import numpy as np
import torch

# ── LoRA for sequence classification (BERT-style) ─────────────────────
MODEL_NAME = 'distilbert-base-uncased'
tokenizer  = AutoTokenizer.from_pretrained(MODEL_NAME)
base_model = AutoModelForSequenceClassification.from_pretrained(
    MODEL_NAME, num_labels=3,
)

# ── Configure LoRA ────────────────────────────────────────────────────
lora_config = LoraConfig(
    task_type=TaskType.SEQ_CLS,  # sequence classification
    r=8,                          # rank — lower = fewer params, higher = more capacity
    lora_alpha=16,                # scaling = alpha/r = 2
    lora_dropout=0.1,             # dropout on LoRA layers
    target_modules=['q_lin', 'v_lin'],  # which layers to add LoRA to
    # For BERT: ['query', 'value']
    # For LLaMA: ['q_proj', 'v_proj', 'k_proj', 'o_proj']
    bias='none',                  # don't train biases
    inference_mode=False,
)

# ── Apply LoRA to model — three lines ─────────────────────────────────
peft_model = get_peft_model(base_model, lora_config)

# Print trainable parameters
peft_model.print_trainable_parameters()
# Output: trainable params: 296,451 || all params: 67,252,227 || trainable%: 0.44%

# ── Verify what is frozen and what is trainable ───────────────────────
print("
Layer-by-layer trainability:")
for name, param in peft_model.named_parameters():
    if param.requires_grad:
        print(f"  TRAINABLE: {name:<60} {param.numel():>10,}")
    # Frozen layers are too many to print — just count them
frozen_count = sum(1 for p in peft_model.parameters() if not p.requires_grad)
print(f"  (+ {frozen_count} frozen parameter tensors)")

# ── Fine-tune with standard Trainer ───────────────────────────────────
LABELS   = ['positive', 'negative', 'neutral']
label2id = {l: i for i, l in enumerate(LABELS)}

reviews = [
    ("Excellent product quality, very happy", 0),
    ("Terrible, broke within a week", 1),
    ("Average product, nothing special", 2),
    ("Amazing value for money", 0),
    ("Worst purchase ever made", 1),
    ("Decent enough for the price", 2),
] * 40

def tokenise(examples):
    return tokenizer(examples['text'], max_length=64,
                     truncation=True, padding=False)

raw   = Dataset.from_dict({
    'text':  [r[0] for r in reviews],
    'label': [r[1] for r in reviews],
})
raw   = raw.train_test_split(test_size=0.2, seed=42)
tok   = raw.map(tokenise, batched=True, remove_columns=['text'])

acc_metric = evaluate.load('accuracy')
def compute_metrics(eval_pred):
    preds = np.argmax(eval_pred.predictions, axis=1)
    return acc_metric.compute(predictions=preds, references=eval_pred.label_ids)

args = TrainingArguments(
    output_dir='./lora-sentiment',
    num_train_epochs=3,
    per_device_train_batch_size=16,
    learning_rate=3e-4,   # LoRA can use higher lr than full fine-tuning
    warmup_ratio=0.1,
    evaluation_strategy='epoch',
    save_strategy='epoch',
    load_best_model_at_end=True,
    seed=42,
)

trainer = Trainer(
    model=peft_model, args=args,
    train_dataset=tok['train'],
    eval_dataset=tok['test'],
    tokenizer=tokenizer,
    data_collator=DataCollatorWithPadding(tokenizer),
    compute_metrics=compute_metrics,
)

print("
Fine-tuning with LoRA (0.44% of parameters):")
trainer.train()
results = trainer.evaluate()
print(f"Accuracy: {results['eval_accuracy']:.4f}")

# QLoRA — Fine-tuning a 7B model on a single GPU
# pip install peft transformers accelerate bitsandbytes

from transformers import (
    AutoTokenizer, AutoModelForCausalLM,
    BitsAndBytesConfig, TrainingArguments,
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
import torch

# ── Step 1: Load model in 4-bit quantisation (QLoRA) ─────────────────
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type='nf4',         # NormalFloat4 — best quality 4-bit
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,    # nested quantisation for extra savings
)

# In production — use a real model:
# model_id = 'meta-llama/Llama-3-8b-hf'  # requires HuggingFace access token
# model_id = 'mistralai/Mistral-7B-v0.3'
# model_id = 'google/gemma-7b'
model_id = 'facebook/opt-125m'   # small model for demonstration

tokenizer = AutoTokenizer.from_pretrained(model_id)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,    # 4-bit quantisation
    device_map='auto',                  # distribute across available GPUs
)

# ── Step 2: Prepare for training ──────────────────────────────────────
# Must call this before adding LoRA — handles gradient checkpointing
model = prepare_model_for_kbit_training(model)

# ── Step 3: Add LoRA ──────────────────────────────────────────────────
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    # Target the attention projection layers
    # For LLaMA/Mistral: ['q_proj', 'v_proj', 'k_proj', 'o_proj', 'gate_proj']
    # For OPT: ['q_proj', 'v_proj']
    target_modules=['q_proj', 'v_proj'],
    lora_dropout=0.05,
    bias='none',
    task_type='CAUSAL_LM',
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

# ── Step 4: Training with gradient checkpointing ──────────────────────
# Gradient checkpointing trades compute for memory —
# recomputes activations during backward instead of storing them
model.gradient_checkpointing_enable()

training_args = TrainingArguments(
    output_dir='./qlora-output',
    num_train_epochs=1,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,    # effective batch = 4 × 4 = 16
    learning_rate=2e-4,               # QLoRA uses higher lr
    fp16=True,
    logging_steps=10,
    optim='paged_adamw_32bit',        # paged optimizer saves memory
    lr_scheduler_type='cosine',
    warmup_ratio=0.03,
    save_strategy='epoch',
)

print("QLoRA configuration:")
print(f"  Rank (r):          {lora_config.r}")
print(f"  Alpha:             {lora_config.lora_alpha}")
print(f"  Target modules:    {lora_config.target_modules}")
print(f"  Quantisation:      4-bit NF4")
print(f"  Effective batch:   {training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps}")

# ── Step 5: Save and load LoRA adapter separately ─────────────────────
# The adapter is tiny — only the LoRA matrices (a few MB)
# model.save_pretrained('./my-lora-adapter')
# tokenizer.save_pretrained('./my-lora-adapter')

# Load later:
# base_model = AutoModelForCausalLM.from_pretrained(model_id, ...)
# model = PeftModel.from_pretrained(base_model, './my-lora-adapter')

from peft import (
    LoraConfig, PromptTuningConfig, PrefixTuningConfig,
    get_peft_model, TaskType, PromptTuningInit,
)
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
import torch

# ── Compare PEFT methods on the same model ────────────────────────────
model_id  = 'facebook/bart-base'   # seq2seq for demonstration
tokenizer = AutoTokenizer.from_pretrained(model_id)

def count_params(model):
    total     = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    return total, trainable

# ── LoRA ──────────────────────────────────────────────────────────────
from transformers import AutoModelForSeq2SeqLM
base = AutoModelForSeq2SeqLM.from_pretrained(model_id)
lora_cfg = LoraConfig(
    task_type=TaskType.SEQ_2_SEQ_LM,
    r=8, lora_alpha=16, lora_dropout=0.05,
    target_modules=['q_proj', 'v_proj'],
    bias='none',
)
lora_model = get_peft_model(base, lora_cfg)
t, tr = count_params(lora_model)
print(f"LoRA:          {tr:>10,} / {t:>10,}  ({tr/t*100:.3f}%)")

# ── Prefix Tuning ─────────────────────────────────────────────────────
base2 = AutoModelForSeq2SeqLM.from_pretrained(model_id)
prefix_cfg = PrefixTuningConfig(
    task_type=TaskType.SEQ_2_SEQ_LM,
    num_virtual_tokens=20,        # 20 trainable prefix tokens
    encoder_hidden_size=768,
)
prefix_model = get_peft_model(base2, prefix_cfg)
t2, tr2 = count_params(prefix_model)
print(f"Prefix tuning: {tr2:>10,} / {t2:>10,}  ({tr2/t2*100:.3f}%)")

# ── Prompt Tuning ─────────────────────────────────────────────────────
base3 = AutoModelForSeq2SeqLM.from_pretrained(model_id)
prompt_cfg = PromptTuningConfig(
    task_type=TaskType.SEQ_2_SEQ_LM,
    num_virtual_tokens=8,
    prompt_tuning_init=PromptTuningInit.TEXT,
    prompt_tuning_init_text="Summarise the following payment dispute: ",
    tokenizer_name_or_path=model_id,
)
prompt_model = get_peft_model(base3, prompt_cfg)
t3, tr3 = count_params(prompt_model)
print(f"Prompt tuning: {tr3:>10,} / {t3:>10,}  ({tr3/t3*100:.3f}%)")

print(f"
Guideline:")
print(f"  Use LoRA for most tasks — best quality per parameter")
print(f"  Use prefix tuning for generation with frequent task switching")
print(f"  Use prompt tuning only for very large models (10B+)")

from peft import PeftModel, LoraConfig, get_peft_model, TaskType
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# ── Simulate: train LoRA then merge for production ────────────────────
MODEL_NAME = 'distilbert-base-uncased'
tokenizer  = AutoTokenizer.from_pretrained(MODEL_NAME)

# 1. Create LoRA model
base_model = AutoModelForSequenceClassification.from_pretrained(
    MODEL_NAME, num_labels=3,
)
lora_config = LoraConfig(
    task_type=TaskType.SEQ_CLS,
    r=8, lora_alpha=16, lora_dropout=0.0,
    target_modules=['q_lin', 'v_lin'],
    bias='none',
)
lora_model = get_peft_model(base_model, lora_config)

# 2. [Training happens here]
# trainer.train()

# 3. Save LoRA adapter (tiny — just A and B matrices)
lora_model.save_pretrained('./my-lora-adapter')
tokenizer.save_pretrained('./my-lora-adapter')

adapter_size = sum(
    p.numel() * 4  # bytes for float32
    for p in lora_model.parameters()
    if p.requires_grad
) / 1024 / 1024
print(f"LoRA adapter size: ~{adapter_size:.1f} MB")

# 4. Merge LoRA into base model — zero inference overhead
merged_model = lora_model.merge_and_unload()
# Now merged_model is a standard nn.Module — no PEFT overhead
# Same weights as if you had done full fine-tuning

print(f"
After merging:")
print(f"  Type:     {type(merged_model).__name__}  ← standard model, no PEFT overhead")
trainable = sum(p.numel() for p in merged_model.parameters() if p.requires_grad)
total     = sum(p.numel() for p in merged_model.parameters())
print(f"  Params:   {total:,} total, {trainable:,} trainable")

# 5. Save merged model
merged_model.save_pretrained('./merged-model')
tokenizer.save_pretrained('./merged-model')

# 6. Load and serve — identical to a standard fine-tuned model
loaded = AutoModelForSequenceClassification.from_pretrained('./merged-model')
loaded.eval()

test_input = tokenizer("Payment declined please help",
                        return_tensors='pt', max_length=64, truncation=True)
with torch.no_grad():
    logits = loaded(**test_input).logits
print(f"
Inference works identically after merge: {logits.shape}")

# ── Serving multiple tasks from one base model ─────────────────────────
print("""
# Production pattern: one base model, multiple LoRA adapters
# Load base once, swap adapters per request

from peft import PeftModel

base = AutoModelForCausalLM.from_pretrained('mistralai/Mistral-7B-v0.3', ...)

# Task 1: customer support
support_model = PeftModel.from_pretrained(base, './lora-support-adapter')

# Task 2: code generation
code_model = PeftModel.from_pretrained(base, './lora-code-adapter')

# Task 3: SQL generation
sql_model = PeftModel.from_pretrained(base, './lora-sql-adapter')

# One GPU, three specialists — much cheaper than three separate 7B models
""")

# vLLM multi-LoRA serving — one base model process, many tenant adapters
# pip install vllm

from vllm import LLM, SamplingParams
from vllm.lora.request import LoRARequest

# ── Load the shared base model once ───────────────────────────────────
llm = LLM(
    model='mistralai/Mistral-7B-v0.3',
    enable_lora=True,
    max_loras=8,            # how many adapters can be resident in GPU memory at once
    max_lora_rank=16,
)

# ── Adapter registry — maps tenant_id to a saved LoRA adapter path ────
TENANT_ADAPTERS = {
    'acme_corp':   './adapters/acme_corp-v3',
    'globex_inc':  './adapters/globex_inc-v1',
    'initech_llc': './adapters/initech_llc-v2',
}

def serve_request(tenant_id: str, prompt: str) -> str:
    adapter_path = TENANT_ADAPTERS.get(tenant_id)
    if adapter_path is None:
        raise ValueError(f"No adapter registered for tenant {tenant_id}")

    # LoRARequest tells vLLM which adapter to apply for this request —
    # vLLM handles loading, caching, and eviction internally
    lora_request = LoRARequest(
        lora_name=tenant_id,
        lora_int_id=hash(tenant_id) % 100000,
        lora_local_path=adapter_path,
    )

    outputs = llm.generate(
        [prompt],
        SamplingParams(temperature=0, max_tokens=200),
        lora_request=lora_request,
    )
    return outputs[0].outputs[0].text

# ── Requests for three different tenants, same GPU, same base model ───
print(serve_request('acme_corp',   "Summarise this support ticket: ..."))
print(serve_request('globex_inc',  "Summarise this support ticket: ..."))
print(serve_request('initech_llc', "Summarise this support ticket: ..."))

# ── Rough cost comparison this architecture is chasing ─────────────────
print("""
1,000 customers, full fine-tuning:
  1,000 x 14GB fp16 models = 14TB of model storage
  Realistically needs its own GPU pool per active customer to avoid reload latency

1,000 customers, shared base + LoRA adapters:
  1 x 14GB base model (loaded once)
  1,000 x ~18MB adapters = ~18GB total adapter storage
  A handful of GPUs serving all tenants, adapters swapped in milliseconds
""")

This suggests LoRA first computes (or could compute) the full-rank weight update and then squeezes it down to rank r — like a JPEG of a photo. That is backwards. LoRA never computes a full-rank update at all; it constrains the trainable update to be low-rank from the very first gradient step, based on the empirical observation (from the original LoRA paper) that the useful update for adapting a pretrained model to a new task tends to live in a low-dimensional subspace. There is nothing to compress because the full-rank version is never formed. If the task genuinely needs a high-rank update, LoRA at low r will underfit — not "lose information" the way compression does.

Rank controls capacity, and more capacity sounds like it should always help — but in practice most tasks saturate well before r=64, and pushing rank higher mostly adds trainable parameters and compute without a measurable accuracy gain, since the intrinsic rank of the needed update is a property of the task, not something you can buy more of. Past that saturation point, a larger r increases the risk of overfitting on small fine-tuning datasets and slows training for no return. The right instinct is to sweep r across 4, 8, 16, and 32 on a validation set for your specific task, not to default to the largest rank you can afford.

During training this is true — the forward pass runs W×x and B@A×x as two separate operations and adds them. But LoRA's defining trick is that B@A is just a matrix, the same shape as W, so it can be added directly into W once training is done: W_merged = W + (α/r)·B@A. After merge_and_unload(), the model has exactly the same architecture and the same number of weight matrices as the original — no adapter, no extra addition, no latency overhead at all. The "extra layer" framing only applies to a few other PEFT methods (adapters, prefix tuning) that genuinely cannot be merged away because they change the computation graph, not just the weights.

LoRA reliably reaches 90–95% of full fine-tuning quality for tasks that adapt the model's style, format, or narrow behaviour — the update genuinely is low-rank there. It is a worse fit when the task requires injecting large amounts of new factual knowledge or substantially shifting the model's underlying capabilities (e.g. teaching a language model an entirely new language it saw little of in pretraining) — that kind of change often needs a higher-rank, more expressive update than LoRA can efficiently represent. Full fine-tuning, continued pretraining, or RAG-for-knowledge remain the better tool in those cases; LoRA did not replace them, it expanded the set of tasks where you no longer need them.

Because each adapter is "just" a small B@A matrix, it is tempting to assume adding two of them (or averaging their weights) gives you a model that is good at both tasks at once. In practice, adapters trained independently occupy overlapping but different subspaces of the weight space, and naively summing them causes destructive interference — the merged behaviour is often worse at both tasks than either adapter alone. Serving multiple tasks from one base model safely means keeping adapters separate and swapping which one is active per request (exactly the pattern PeftModel.from_pretrained shown earlier supports), or using purpose-built multi-adapter composition methods, not ad hoc addition.

The hypothesis, backed empirically by the original LoRA paper (and earlier work on intrinsic dimensionality of fine-tuning), is that the change a pretrained model needs to specialise for a downstream task lives in a much lower-dimensional subspace than the full weight matrix it is applied to — even though the base weight itself is high-rank and uses its full capacity to encode broad pretraining knowledge. Concretely: adapting a 768×768 attention matrix to a new task does not require exploring all 589,824 independent directions of change; a rank-8 or rank-16 subspace captures nearly all of the useful signal. LoRA exploits this directly by parameterising the update as B@A with rank r, training only 2×r×d parameters instead of d², and getting comparable task performance because it was never necessary to search the full-rank space in the first place.

Rank r sets the dimensionality of the low-rank update — it is the actual capacity knob: higher r means more trainable parameters and a richer space of representable updates, at the cost of more compute and overfitting risk on small datasets. Alpha is a separate scaling factor applied to the LoRA output as α/r, controlling how strongly the update perturbs the frozen base weights relative to its own magnitude — it does not add capacity, it tunes how aggressively that capacity is expressed. In practice most teams do not tune them independently: a common heuristic is alpha = 2×r (so the effective scale stays roughly constant as rank changes), then sweep r itself — typically 8 or 16 for classification-style tasks and up to 32-64 for more complex generation tasks — against a validation metric.

The decision axis is how much you need to change the model versus how much compute and data you have. Full fine-tuning updates every weight — highest ceiling on quality and the only real option when the task requires deep, broad changes to the model's behaviour or knowledge, but it needs the most GPU memory (model + gradients + optimiser states for every parameter) and the most labelled data to avoid overfitting all those parameters. LoRA sits in the middle: 90-95% of full fine-tuning quality at a fraction of the memory, because it exploits the low-rank structure of most task adaptations, and it is the default choice for the vast majority of production fine-tuning today. Prompt tuning goes further — training only a handful of soft input tokens with no per-layer weight changes at all — which underperforms LoRA at small-to-medium model scale but becomes surprisingly competitive at 10B+ parameters, where the frozen base model is already so capable that a small nudge to its input is enough to steer behaviour.

The forward pass during training is h = W×x + (α/r)·(B@A)×x, where W is frozen and A, B are the only trainable matrices. Because matrix multiplication distributes over addition, this is algebraically identical to h = (W + (α/r)·B@A)×x — a single matrix, the same shape as W, multiplied once against x. So after training finishes, you compute W_merged = W + (α/r)·B@A exactly once, write it back in place of W, and discard A and B entirely. The resulting model has the identical architecture, weight shapes, and forward pass as a model that was fully fine-tuned from scratch — there is no separate adapter branch left to execute, so inference cost is exactly the same as the unmodified base model. This only works for LoRA (and similar linear reparameterisations) specifically because the adaptation is additive and linear in x; adapters and prefix tuning change the computation graph itself and cannot be collapsed this way.

QLoRA combines LoRA with 4-bit quantisation of the frozen base model weights (NF4 — NormalFloat4 — plus double quantisation of the quantisation constants themselves), which shrinks a 7B model's weight storage from roughly 14-28GB down to about 3.5GB. Because only the small LoRA matrices are trained, gradients and optimiser states are computed and stored only for those — a fraction of a percent of total parameters — so the two most memory-hungry pieces of full fine-tuning (per-parameter gradients and per-parameter Adam state) barely register. The catch is that 4-bit quantisation is inherently lossy: it introduces small numerical errors into every frozen weight, which in principle could hurt output quality. In practice, the QLoRA paper showed this loss is largely absorbed because the LoRA adapter is trained on top of the quantised weights and learns to compensate — but it does mean QLoRA models can be marginally less precise on tasks sensitive to exact numerical behaviour, and it adds a dequantisation cost at inference unless you also merge and requantise afterward.