LLMs · Fine-tuning with LoRA

LoRA freezes the pretrained weights and learns a low-rank update BA; a LoRA layer written from scratch, then PEFT applied to GPT-2 to count how few parameters are trained.

Full fine-tuning updates every weight, which is expensive and gives you a full model copy per task. LoRA (Hu et al., 2021):

freeze the pretrained weight W
learn a low-rank update: W’ = W + (α/r)·B·A, with A ∈ ℝ^{r×d}, B ∈ ℝ^{d×r} and a small r (e.g. 8)
initialise B = 0, so training starts exactly from the pretrained model

Trainable numbers per matrix drop from d² to 2dr. After training, BA can be merged into W, so inference has no extra cost.

import torch, torch.nn as nn
from transformers import AutoModelForCausalLM
from peft import LoraConfig, get_peft_model
torch.manual_seed(0)

class LoRALinear(nn.Module):
    def __init__(self, base: nn.Linear, r=8, alpha=16):
        super().__init__()
        self.base = base.requires_grad_(False)
        self.A = nn.Parameter(torch.randn(r, base.in_features) * 0.01)
        self.B = nn.Parameter(torch.zeros(base.out_features, r))
        self.scale = alpha / r
    def forward(self, x):
        return self.base(x) + (x @ self.A.T @ self.B.T) * self.scale

base = nn.Linear(768, 768)
lora = LoRALinear(base)
x = torch.randn(4, 768)
print("unchanged at init (B=0):", torch.allclose(lora(x), base(x)))
train = sum(p.numel() for p in lora.parameters() if p.requires_grad)
print(f"trainable {train:,} vs full {768 * 768 + 768:,} ({train / (768 * 768 + 768):.2%})")

lora.B.data.normal_()                                    # pretend we trained
merged = base.weight + lora.scale * (lora.B @ lora.A)
print("merged weight gives same output:", torch.allclose(lora(x), x @ merged.T + base.bias, atol=1e-4))

gpt2 = AutoModelForCausalLM.from_pretrained("gpt2")
peft_model = get_peft_model(gpt2, LoraConfig(r=8, lora_alpha=16, target_modules=["c_attn"], lora_dropout=0.0, task_type="CAUSAL_LM"))
peft_model.print_trainable_parameters()

Output:

unchanged at init (B=0): True
trainable 12,288 vs full 590,592 (2.08%)
merged weight gives same output: True
trainable params: 294,912 || all params: 124,734,720 || trainable%: 0.2364

LLMs · Fine-tuning with LoRA

Sources