Supervised fine-tuning

Fine-tuning the base model

Now the real thing. The cell below downloads tiny-gpt, the base model from the last lesson. It fine-tunes it on 300 chats, one per product in the catalog. Each chat asks the model to describe a product. The answer is the first sentence of that product's description.

Then it tests the model on 20 products that aren't in the catalog. Each one pairs a name the model knows with a kind of product it never saw under that name. An answer passes if it's one sentence, names the right product, and ends with <end>. The cell scores the model before and after training.

The comments split the cell into three parts.

1. The data. Build one chat per product, with the prompt labeled -100. Chats have different lengths, so shorter ones get extra <end> tokens to bring them all to one length. That padding is labeled -100 too, so it doesn't count.
2. The test. Pick the unseen products and write a function that scores the model's answers.
3. Fine-tune. This is the training loop you know. Pick 16 chats at random, work out the masked loss, and nudge the weights. Repeat for about 300 steps, printing the loss every 100.

It takes a while to run. The Run button counts the seconds.

import re
import requests
import torch
import torch.nn as nn
from tinygpt import load_pretrained

model, tok = load_pretrained("tiny-gpt")                      # the base model
end = tok.token_id("<end>")
docs = requests.get("https://api.gradient.test/v1/corpus").json()["text"].strip().split("\n")
product = re.compile(r"^(?:Meet the|Our|The) ([A-Z][a-z]+) (.+?)(?:,| is| are) made")

# 1. the data: one chat per catalog product, answered with the description's first sentence
rows = []
for d in docs:
    name, kind = product.match(d).groups()
    prompt = tok.encode(f"<user> Describe the {name} {kind} . <assistant>")
    answer = tok.encode(d.split(". ")[0] + ".") + [end]
    rows.append((prompt + answer, [-100] * len(prompt) + answer))     # learn the answer only
T = max(len(ids) for ids, _ in rows)
X = torch.tensor([ids + [end] * (T - len(ids)) for ids, _ in rows])  # pad to one length
Y = torch.tensor([lab + [-100] * (T - len(lab)) for _, lab in rows])
print(f"{len(rows)} training chats, up to {T} tokens each")

# 2. the test: products that aren't in the catalog
seen = {product.match(d).groups() for d in docs}
names, kinds = sorted({n for n, _ in seen}), sorted({k for _, k in seen})
unseen = [(n, k) for n in names for k in kinds if (n, k) not in seen][::8]

@torch.no_grad()
def evaluate(model):
    good, replies = 0, []
    for n, k in unseen:
        ids = torch.tensor([tok.encode(f"<user> Describe the {n} {k} . <assistant>")])
        reply = tok.decode(model.generate(ids, 30, temperature=0.7, stop_id=end)[0, ids.shape[1]:].tolist())
        good += f"{n} {k}" in reply and reply.endswith("<end>") and reply.count(".") == 1
        replies.append(f"{n} {k}: {reply}")
    return good / len(unseen), replies

torch.manual_seed(0)
score, replies = evaluate(model)
print(f"\nbefore: {score:.0%} of {len(unseen)} new products get a one-sentence answer that names them")
print("  ", replies[0])

# 3. fine-tune: ordinary next-token training, with the prompt positions ignored
model.train()
opt = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.01)
torch.manual_seed(1)
for step in range(301):
    b = torch.randint(0, len(X), (16,))
    logits = model(X[b][:, :-1])
    loss = nn.functional.cross_entropy(logits.reshape(-1, logits.shape[-1]), Y[b][:, 1:].reshape(-1))
    opt.zero_grad()
    loss.backward()
    opt.step()
    if step % 100 == 0:
        print(f"step {step:>3}  loss {loss.item():.3f}")

model.eval()
torch.manual_seed(0)
score, replies = evaluate(model)
print(f"\nafter: {score:.0%}")
for r in replies[:4]:
    print("  ", r)

Before fine-tuning, the base model treats the chat as more catalog text and rambles on. After about 300 steps of 16 chats, it has seen 4,800 chats. That's about 16 passes over the 300 chats. A tiny model needs many more passes than a large one.

Now it answers in the new format for most new products: one sentence that stops. That includes products it never saw described.

The results also show the limits of SFT. Asked about the Aspen fleece, the model describes an "Aspen headlamp, made from a light fleece". It mixed up the product it was asked about with a familiar one, and it invented a material. SFT taught the format. The knowledge, and its gaps, came from pretraining. (torch here is the course's in-browser PyTorch subset.)

Getting SFT right

Data matters most. A few hundred to a few thousand good, consistent demonstrations often beat tens of thousands of mediocre ones. Every example teaches something. If a tenth of your examples are sloppy, the model learns to be sloppy a tenth of the time. Cover the range of requests you expect. Include ones the model should decline or answer briefly.

Watch for overfitting. Overfitting means memorizing the training examples instead of learning the general pattern. Fine-tuning datasets are small, so models memorize them fast. Hold out some examples that you don't train on, and track the loss on them. This is the validation loss. Stop when it starts to rise, which is typically after one to three passes over the data. An overfit model repeats training answers word for word and gets worse on anything new.

Use a small learning rate for big models. The learning rate sets how big each training nudge is. Full fine-tuning of a large model typically uses a learning rate around 0.00001, written $10^{-5}$. That's big enough to change behavior, but small enough not to wreck what pretraining learned. Our tiny model needs a much larger one to change in a few hundred steps. The cell used lr=3e-3, which is Python's way of writing 0.003.

In plain words: $10^{-5}$ means 1 divided by 10 five times, which is 0.00001. Python writes it as 1e-05. The tiny model's rate, 0.003, is 0.003 ÷ 0.00001 = 300 times bigger.

print(10 ** -5)              # Python's short way of writing 0.00001
print(round(3e-3 / 1e-5))    # how many times bigger the tiny model's rate is

The first line prints 1e-05, and the second prints 300.

Mind what you might break. Training only on your narrow task can wear away other skills. This is called catastrophic forgetting. Two things limit the damage. You can mix some general chat data into the fine-tuning set. Or you can use parameter-efficient methods like LoRA, which come later in this phase.

Test on held-out prompts that look like real use, as the cell did with products outside the catalog. Compare the result with the base model. Also compare it with plain prompting of a stronger model, which is often good enough.

With real tools

On your own computer or a rented GPU, the usual tool is Hugging Face's TRL library. Its SFTTrainer takes a dataset of chats in the standard messages format. It applies the model's chat template for you. Then it runs the same loop you just ran:

from datasets import load_dataset
from trl import SFTConfig, SFTTrainer

dataset = load_dataset("json", data_files="brand_chats.jsonl", split="train")   # {"messages": [...]} per line
trainer = SFTTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",            # any causal LM on the Hugging Face Hub
    train_dataset=dataset,
    args=SFTConfig(output_dir="brand-model", num_train_epochs=2, learning_rate=2e-5),
)
trainer.train()

A few settings are worth reading. The file brand_chats.jsonl is a JSONL file, which means one JSON record per line, here one chat per line. num_train_epochs=2 asks for two passes over the data. (One full pass is called an epoch.) learning_rate=2e-5 is 0.00002, in the small range big models need.

Hosted APIs offer the same thing without managing hardware. Some model providers let you upload a JSONL file of chats. You get back a fine-tuned version of one of their models.

For marketers: before fine-tuning for brand voice, try a strong model first. Give it a style guide and five good examples in the prompt. Fine-tune when that isn't consistent enough at your volume, or when a small fine-tuned model would cut costs. Then start with a few hundred excellent examples.