Part 2 · Large language models · module 3 of 4 · 21 min

Fine-tune pretrained models

Adapting an existing model to your data or task — full and partial fine-tuning, freezing, PEFT and LoRA, classification heads, clipping and schedulers.

The essence

Fine-tuning is pretraining continued on a smaller, narrower dataset with a lower learning rate. Nothing about the mechanism changes — same forward pass, same loss, same optimiser. What changes is that the model starts competent, so you are nudging existing structure rather than building it, which is why it takes minutes rather than months.

Why it matters

This is the part you will actually do. It is also where the characteristic failure modes live: catastrophic forgetting from too high a learning rate, overfitting on a small dataset, and the choice between full fine-tuning and a parameter-efficient method that determines whether the job fits on your GPU at all.

Topics covered (21)

01What fine-tuning means
02Fine-tuning a pretrained GPT-2 on a single text
03The effect of learning rate when fine-tuning
04Generating text from pretrained models
05Steering output toward a target with fine-tuning
06Two models, two authorial styles
07Quantifying the effect of style fine-tuning
08Making two fine-tuned models converse
09Partial fine-tuning by freezing attention weights
10Targeted freezing: which layers to unfreeze
11Parameter-efficient fine-tuning (PEFT) and LoRA
12Fine-tuning a code-completion model
13Fine-tuning a code model on a narrow domain
14Fine-tuning BERT for classification
15Sentiment analysis with BERT
16Gradient clipping and learning-rate schedulers
17Combining clipping, freezing and scheduling
18Saving and loading trained models
19Using a classifier to detect style
20Tracking how a fine-tuned model evolves over training
21When fine-tuning is the wrong tool

What changes, and what does not

The training loop is identical to pretraining. The differences are all in the settings: a learning rate around 1e-5 to 5e-5 instead of 3e-4, a handful of epochs instead of one pass over a vast corpus, and often only some parameters unfrozen.

The reason the learning rate must drop is that the pretrained weights are a good solution you are trying not to destroy. Large steps walk out of that basin, and the model forgets general language ability while learning your narrow data — catastrophic forgetting, which appears as fluent output that has lost all knowledge outside the fine-tuning set.

Fine-tuning GPT-2 on one text

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from torch.utils.data import DataLoader, Dataset

tok = AutoTokenizer.from_pretrained("gpt2")
tok.pad_token = tok.eos_token              # GPT-2 ships without a pad token
model = AutoModelForCausalLM.from_pretrained("gpt2")

class Chunks(Dataset):
    def __init__(self, text, block=128):
        ids = tok(text, return_tensors="pt")["input_ids"][0]
        self.chunks = ids.split(block)[:-1]      # drop the ragged tail
    def __len__(self):  return len(self.chunks)
    def __getitem__(self, i): return self.chunks[i]

loader = DataLoader(Chunks(open("gulliver.txt").read()), batch_size=4, shuffle=True)
opt = torch.optim.AdamW(model.parameters(), lr=5e-5)     # note the small lr

model.train()
for epoch in range(3):
    for batch in loader:
        # For causal LM, labels == input_ids; the model shifts them internally.
        out = model(input_ids=batch, labels=batch)
        opt.zero_grad(set_to_none=True)
        out.loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
    print(epoch, out.loss.item())

Note: Hugging Face causal models shift labels for you. Shifting them yourself as well trains the model to predict the token it was just given.

Learning rate, and what it does to the model

Learning rate	Outcome
1e-6	Barely moves; style unchanged after a few epochs
1e-5 – 5e-5	The useful band: style adapts, general ability survives
1e-4	Fast adaptation, visible degradation elsewhere
1e-3 and above	Catastrophic forgetting — fluent, repetitive, knowledge gone

Generating from a fine-tuned model

model.eval()
prompt = tok("In the morning I", return_tensors="pt")

out = model.generate(
    **prompt,
    max_new_tokens=60,
    do_sample=True,          # False makes every other sampling arg inert
    temperature=0.8,
    top_p=0.9,
    repetition_penalty=1.1,
    no_repeat_ngram_size=3,
    pad_token_id=tok.eos_token_id,
)
print(tok.decode(out[0], skip_special_tokens=True))

# Quantifying a style change without reading output: compare the loss the
# fine-tuned model assigns to held-out text in the target style against a
# control style. A larger gap than the base model means the style was learned
# rather than the text memorised.

Note: do_sample defaults to False. Setting temperature and top_p while leaving it False changes nothing — a very common source of 'why is my output identical every time'.

Partial fine-tuning: freezing

You need not train everything. Setting requires_grad = False on a parameter excludes it from the update, which cuts memory (no gradients, no optimiser state) and constrains what can change.

The general finding is that early layers encode broadly useful structure — tokenization-adjacent regularities, basic syntax — while later layers encode task and style specifics. So freezing early layers and tuning later ones adapts the model with far less risk of forgetting.

Freezing selectively is also an experiment. Freeze all attention and tune only the MLPs, and any style change you observe must have been implementable in per-position computation alone. Freezing is a causal probe as well as an efficiency measure.

Freezing, three ways

# 1. Freeze everything, then unfreeze the last two blocks and the final norm.
for p in model.parameters():
    p.requires_grad = False
for block in model.transformer.h[-2:]:
    for p in block.parameters():
        p.requires_grad = True
for p in model.transformer.ln_f.parameters():
    p.requires_grad = True

# 2. Freeze only the attention projections, leaving the MLPs trainable.
for block in model.transformer.h:
    for p in block.attn.parameters():
        p.requires_grad = False

# 3. Freeze the embeddings — often sensible, and they are the largest tensor.
model.transformer.wte.weight.requires_grad = False

trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total     = sum(p.numel() for p in model.parameters())
print(f"{trainable:,} / {total:,}  ({100 * trainable / total:.1f}%)")

# Build the optimiser AFTER freezing, over the trainable parameters only.
opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=5e-5)

Parameter-efficient fine-tuning and LoRA

Full fine-tuning of a 7B model needs the weights, their gradients, and two AdamW moment buffers — roughly four times the model in memory, before activations. That is out of reach on ordinary hardware.

LoRA's observation is that the update a fine-tune applies is effectively low-rank. So instead of learning a full ΔW, learn two thin matrices B (d×r) and A (r×k) with r around 8 to 32, and use W + BA. The base weights stay frozen; only A and B are trained, typically well under 1% of the parameters.

A is initialised randomly and B to zeros, so BA starts at exactly zero and the model begins as the unmodified base. After training the product can be folded into W, meaning zero inference overhead. QLoRA goes further by quantising the frozen base to 4-bit, which is what makes single-GPU fine-tuning of large models routine.

LoRA, from scratch and with the library

import torch, torch.nn as nn

class LoRALinear(nn.Module):
    """Wraps a frozen Linear with a trainable low-rank update."""
    def __init__(self, base: nn.Linear, r=8, alpha=16):
        super().__init__()
        self.base = base
        for p in self.base.parameters():
            p.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))   # starts at 0
        self.scale = alpha / r

    def forward(self, x):
        return self.base(x) + self.scale * (x @ self.A.T @ self.B.T)

# In practice, use peft:
from peft import LoraConfig, get_peft_model

config = LoraConfig(
    r=16, lora_alpha=32, lora_dropout=0.05,
    target_modules=["c_attn"],       # which projections get adapters
    task_type="CAUSAL_LM",
)
model = get_peft_model(base_model, config)
model.print_trainable_parameters()   # e.g. 0.24% of all parameters

Fine-tuning BERT for classification

A classification fine-tune replaces the language-modelling head with a small classifier on top of the [CLS] token's final hidden state, and trains with cross-entropy over class labels. The encoder learns to concentrate whatever distinguishes the classes into that one position.

AutoModelForSequenceClassification does the head-swapping for you. The warning about newly initialised weights is expected and correct — that head is random and is what you are training.

BERT sentiment classification

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

tok = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModelForSequenceClassification.from_pretrained(
    "bert-base-uncased", num_labels=2)      # head is randomly initialised

batch = tok(texts, truncation=True, padding=True, max_length=256, return_tensors="pt")
labels = torch.tensor(ys)

opt = torch.optim.AdamW(model.parameters(), lr=2e-5)   # 2e-5 is the BERT default

model.train()
out = model(**batch, labels=labels)
out.loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()

model.eval()
with torch.no_grad():
    preds = model(**batch).logits.argmax(dim=-1)

# A trained classifier is also a measurement instrument: train one to tell two
# authors apart, then score a model fine-tuned on one of them. Classifier
# confidence over training checkpoints quantifies how fast the style was
# acquired, without anyone reading the samples.

Gradient clipping and learning-rate schedules

Gradient clipping rescales the whole gradient vector when its norm exceeds a threshold, preserving direction and capping magnitude. One pathological batch can otherwise produce an enormous step that wrecks a model that was training well. Clipping the global norm to 1.0 is close to free and near-universal.

A schedule shapes the learning rate over training. Warm-up matters most: at step zero the optimiser's variance estimates are meaningless, and a full-size step is a shot in the dark — hence a linear ramp over the first few hundred steps. After that, cosine decay to near zero lets the model settle into a minimum rather than bouncing around it.

Warm-up plus cosine decay

import math, torch
from torch.optim.lr_scheduler import LambdaLR

def warmup_cosine(optimizer, warmup_steps, total_steps, min_ratio=0.1):
    def fn(step):
        if step < warmup_steps:
            return step / max(1, warmup_steps)              # linear ramp
        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
        return min_ratio + (1 - min_ratio) * 0.5 * (1 + math.cos(math.pi * progress))
    return LambdaLR(optimizer, fn)

sched = warmup_cosine(opt, warmup_steps=100, total_steps=len(loader) * epochs)

for epoch in range(epochs):
    for batch in loader:
        loss = model(**batch, labels=batch["input_ids"]).loss
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        sched.step()          # per batch, not per epoch

Note: Calling sched.step() once per epoch instead of per batch is a frequent bug: the schedule then completes in a handful of steps and the rate sits at its floor for the rest of training.

Saving and loading

# Preferred: the library format. Portable, and reloads config and tokenizer.
model.save_pretrained("./my-model")
tok.save_pretrained("./my-model")
model = AutoModelForCausalLM.from_pretrained("./my-model")

# A training checkpoint needs optimiser and scheduler state too, or resuming
# restarts Adam's moment estimates from zero and the loss spikes.
torch.save({
    "model": model.state_dict(),
    "optimizer": opt.state_dict(),
    "scheduler": sched.state_dict(),
    "epoch": epoch,
}, "checkpoint.pt")

ckpt = torch.load("checkpoint.pt", map_location="cpu")
model.load_state_dict(ckpt["model"])
opt.load_state_dict(ckpt["optimizer"])

# For LoRA, save only the adapter — a few megabytes instead of gigabytes.
model.save_pretrained("./adapter")   # peft writes just A and B

Worth remembering

Fine-tuning uses a learning rate 10–100× smaller than pretraining; too high erases the pretrained knowledge.
LoRA trains small low-rank adapters and leaves the base weights frozen — typically under 1% of parameters.
Freeze early layers and tune later ones: early layers encode general structure, later layers encode task specifics.