Fine-tune an open model

Step 3: curate the training data

To curate data means to choose it with care, the way a museum curator chooses what goes on show. Fewer, cleaner examples beat more, noisier ones.

The preference dataset from the preference-data and RLHF lessons gives you a head start. For each prompt, the model wrote four answers, and the dataset kept the best of the four as the "chosen" answer. But the best of four can still break a guideline. So the project adds one more filter. It keeps a chosen answer only if the checker found no flaws at all, and only if it's about a training product.

This trick has a name: rejection-sampling fine-tuning. You sample several answers, reject the ones a checker fails, and fine-tune on the rest. It's a cheap way to get clean examples.

Steps 4 and 5: baseline, training, test

A baseline is the score before you change anything. Without one, you can't tell whether training helped.

A checkpoint is a saved copy of the model's weights at one moment in training. It's like a save point in a video game. The project trains for 120 steps. Every 40 steps, it scores the model on the validation products. Whenever the validation score improves, it saves a checkpoint. At the end, it loads the best checkpoint and scores the test products, once.

The cell runs all five steps. It takes a while, and the Run button counts the seconds. (torch here is the course's in-browser PyTorch subset.)

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

model, tok = load_pretrained("tiny-gpt-chat")
end = tok.token_id("<end>")
pairs = requests.get("https://api.gradient.test/v1/preferences").json()["data"]
kinds = sorted({" ".join(p["prompt"].split()[3:-1]) for p in pairs})
test_names = {"Kestrel", "Lark", "River", "Sierra", "Summit", "Willow"}
val_names = {"Alder", "Basin", "Cove", "Ember"}

# 1. the metric: violations of the brand guidelines on products the model never trained on
@torch.no_grad()
def violations(model, names):
    torch.manual_seed(0)
    items = []
    for product in [f"{n} {k}" for n in sorted(names) for k in kinds]:
        ids = torch.tensor([tok.encode(f"<user> Describe the {product} . <assistant>")])
        reply = tok.decode(model.generate(ids, 60, temperature=0.8, stop_id=end)[0, ids.shape[1]:].tolist())
        items.append({"product": product, "text": reply.split("<end>")[0].strip()})
    flaws = [r["flaws"] for r in requests.post("https://api.gradient.test/v1/brand-check", json={"items": items}).json()["results"]]
    return sum(bool(f) for f in flaws) / len(flaws), sum(len(f) for f in flaws) / len(flaws)

# 2-3. curated training data: answers that pass every guideline, for training products only
good = [p for p in pairs if not p["chosen_flaws"] and p["prompt"].split()[2] not in test_names | val_names]
rows = []
for p in good:
    prompt, answer = tok.encode(f"<user> {p['prompt']} <assistant>"), tok.encode(p["chosen"]) + [end]
    rows.append(((prompt + answer)[:65], ([-100] * len(prompt) + answer)[:65]))
T = max(len(ids) for ids, _ in rows)
X = torch.tensor([ids + [end] * (T - len(ids)) for ids, _ in rows])
Y = torch.tensor([lab + [-100] * (T - len(lab)) for _, lab in rows])
print(f"{len(rows)} curated training examples")

# 4. baseline, then training with a validation check every 40 steps
test_before = violations(model, test_names)
best = (violations(model, val_names), 0, copy.deepcopy(model.state_dict()))
print(f"step   0: validation {best[0][0]:.0%} of answers break a guideline ({best[0][1]:.2f} per answer)")
model.train()
opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
torch.manual_seed(1)
for step in range(1, 121):
    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 % 40 == 0:
        model.eval()
        val = violations(model, val_names)
        model.train()
        print(f"step {step:>3}: validation {val[0]:.0%} of answers break a guideline ({val[1]:.2f} per answer)   training loss {loss.item():.3f}")
        if val[1] < best[0][1]:
            best = (val, step, copy.deepcopy(model.state_dict()))

# 5. keep the best checkpoint and test once
model.load_state_dict(best[2])
model.eval()
test_after = violations(model, test_names)
print(f"\nbest checkpoint: step {best[1]}")
print(f"test, before: {test_before[0]:.0%} of answers break a guideline ({test_before[1]:.2f} per answer)")
print(f"test, after:  {test_after[0]:.0%} of answers break a guideline ({test_after[1]:.2f} per answer)")

Reading the log

Read the log the way a practitioner would.

• Only 61 clean examples survived curation, from 293 pairs, because the checker is strict. Even so, 61 examples were enough to change behavior. That echoes the LIMA finding from the base-vs-chat lesson: a small set of carefully chosen examples can go a long way.
• Training loss fell the whole time, but validation didn't. Validation improved up to step 80 and got worse by step 120. The model was starting to memorize its 61 examples. Without the validation set, you'd have shipped the step-120 model.
• The test result is the honest number. It was measured once, on products never used for training or for choosing the checkpoint. Violations per answer fell from 1.44 to 0.92, about a third fewer. The share of answers with any violation fell from 85% to 53%.

The DPO lesson reached a similar improvement with a different method. When you have clean examples, plain SFT on curated data is often the simpler choice. Preference methods earn their extra complexity when you have comparisons ("this answer beats that one") rather than perfect examples.

The same project on a real open model

On your own GPU, or a free notebook GPU from Google Colab, the workflow is the same. Only the tools change.

1. Save the curated chats as JSON lines in the messages format. That's one conversation per line, written as a list of messages, each with a role (user or assistant) and its text.
2. Split by product into train, validation and test files.
3. Fine-tune a small open model with LoRA, from the LoRA lesson.

# pip install transformers trl peft datasets
from datasets import load_dataset
from peft import LoraConfig
from trl import SFTConfig, SFTTrainer

data = load_dataset("json", data_files={"train": "brand_train.jsonl", "validation": "brand_val.jsonl"})
trainer = SFTTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    train_dataset=data["train"],
    eval_dataset=data["validation"],
    peft_config=LoraConfig(r=16, lora_alpha=32, target_modules="all-linear", task_type="CAUSAL_LM"),
    args=SFTConfig(output_dir="brand-lora", num_train_epochs=3, learning_rate=2e-4,
                   eval_strategy="epoch", save_strategy="epoch", load_best_model_at_end=True),
)
trainer.train()
trainer.save_model("brand-lora")       # saves just the adapter, a few megabytes

Here's what the main settings mean:

• model names a small open chat model with about half a billion weights.
• LoraConfig trains small LoRA updates instead of every weight. r=16 is the rank, the size of each small update. target_modules="all-linear" adds an update to the linear layers in every block of the model.
• num_train_epochs=3 means three passes over the training data. The learning rate 2e-4 is 0.0002.
• eval_strategy and save_strategy set to "epoch" mean: after each pass, measure the validation loss and save a checkpoint.
• save_model saves only the adapter, the small LoRA update, which is a few megabytes.

load_best_model_at_end does what the cell did by hand. It keeps the checkpoint with the best validation score. But that score is the validation loss, and loss is only a stand-in for what you care about. A low loss means the model finds the validation answers likely. It doesn't prove the model's own answers follow the guidelines. So also score each saved checkpoint with your real metric on the validation prompts. That could be a guideline checker or an LLM judge. Then run the chosen checkpoint once on the test prompts.

A checklist for real projects

1. Write the metric first. Make it automatic if you can: a checker, a classifier, or an LLM judge with a rubric whose scores you've checked against people's.
2. Get a baseline from a well-prompted model, including a larger hosted model. If prompting meets the bar, stop there.
3. Split by the thing you'll see new examples of, such as new products, new customers or new months. A split like this, not a random one, makes the test reflect real use.
4. Curate. Fewer, cleaner examples beat more, noisier ones. Read a sample of your data.
5. Watch validation during training and keep the best checkpoint. Expect a small dataset to overfit within a few passes.
6. Test once. Compare against the baseline on every criterion you care about, and read some outputs yourself.
7. Plan for upkeep. Retrain when the base model, the guidelines or the inputs change.

For marketers: this checklist works as a brief for an agency or an internal team proposing a fine-tune. Ask to see the metric, the baseline, the split, and the test result before and after.