Fine-Tuning with PEFT — LoRA and Adapters

A senior engineer at Razorpay 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
""")