Project 4: Fine-tune a model with LoRA
Pre-train a small character model on one text, fine-tune it with LoRA on a different style, compare rank 1/4/16 and full fine-tuning on perplexity, trained parameters and catastrophic forgetting — and write the conclusion in the lab journal.
Theory
The base model only knows what it was pre-trained on. LoRA adds BA to the attention projections (and optionally the MLP) and trains only those. Measure three things: perplexity on the new style, parameters trained, and perplexity on the old text afterwards (forgetting).
Sub-tasks
- LoRA injection —
inject_lora(model, r, alpha, targets)replaces the given nn.Linear modules (names containing 'qkv'/'proj'/'mlp') with LoRALinear and freezes the rest;count_trainable(model). - perplexity —
perplexity(model, data, ctx)= exp(mean cross-entropy) over sliding windows. - the experiment —
experiment(ranks, steps, seed)→ dict per configuration {ppl_new, ppl_old, trainable}, including 'full'.
Passes when: ppl_gain_r4 >= 1.2
The starter code
runs in an isolated sandbox on the serverimport math
import torch
import torch.nn as nn
import torch.nn.functional as F
from tinygpt import CharTokenizer, TinyGPT, train, batches # levereras i readonly-filen tinygpt.py
class LoRALinear(nn.Module):
def __init__(self, base, r=4, alpha=8):
super().__init__()
self.base = base
for p in base.parameters():
p.requires_grad = False
self.r, self.alpha = r, alpha
self.A = nn.Parameter(torch.randn(r, base.in_features) * 0.01)
self.B = nn.Parameter(torch.zeros(base.out_features, r))
def forward(self, x):
return self.base(x) + (self.alpha / self.r) * (x @ self.A.T) @ self.B.T
def inject_lora(model, r=4, alpha=8, targets=("qkv", "proj")):
# TODO: frys allt; för varje named_module som är nn.Linear och vars namn innehåller något target: byt ut mot LoRALinear
...
def count_trainable(model):
return sum(p.numel() for p in model.parameters() if p.requires_grad)
@torch.no_grad()
def perplexity(model, data, ctx=64, n=40, seed=0):
# TODO: n slumpade fönster (generator med seed); cross-entropy; return exp(medel)
...
def experiment(ranks=(1, 4, 16), steps=300, seed=0):
# TODO: förträna på base_text (600 steg), spara; för varje r: kopiera, inject_lora, träna på new_text; mät ppl_new/ppl_old/trainable; 'full' = träna alla
...
You write the code; tests you cannot see decide whether it holds up. Create a free account to run the lab.
Try the diagnosticCreate a free accountExpected results
LoRA r=4 lowers the perplexity on the new style by at least 20 % compared with the base model, trains < 15 % of the parameters, and forgets less than full fine-tuning (lower ppl_old).
Common mistakes
- Injects LoRA into the embeddings/head as well — then too much is trained.
- Forgets to freeze the base → 'LoRA' becomes full fine-tuning.
- Measures perplexity on the training windows.