Skip to content
AI-grafen
FAI engineeringProject· about 180 min· server sandbox

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

  1. 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).
  2. perplexity — perplexity(model, data, ctx) = exp(mean cross-entropy) over sliding windows.
  3. 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 server
import 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 account

Expected 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.