Supervised Fine-Tuning Tutorial: From Base LLM to Chat Model

A new cashier joins a bank branch on MG Road. He already speaks fluent Hindi and English, reads well and knows arithmetic.

What he doesn't know is the counter routine: greet the customer, confirm the account number, answer in one or two sentences, and close with "Anything else?" His manager doesn't teach him language again. She shows him a few hundred real counter conversations, and he copies the pattern.

That is supervised fine-tuning. The cashier's fluency comes from pretraining. The counter conversations are the SFT dataset. Copying only the cashier's lines, not the customer's, is loss masking. And the fixed routine of greet, answer, close is the chat template.

What is supervised fine-tuning?

Supervised fine-tuning (SFT) is the training stage that takes a pretrained base language model and trains it further on curated prompt-and-response conversations, using the same next-token loss but only on the response tokens, so the model learns to follow instructions in a chat format.

This is stage 6 of the build-an-LLM roadmap. It starts from the checkpoint you produced in pretraining an LLM from scratch and reads the JSON Lines chat data from data preparation. For how SFT compares with preference and reinforcement methods, see SFT vs RLHF vs GRPO.

Same model, same loss function. Only the data and the mask change.

How does a base model differ from a chat model?

Ask your pretrained 64M base model "What is the capital of Madhya Pradesh?" and it will probably continue the text: another question, a list of quiz items, or a paragraph that drifts off. It has never seen the pattern "question, then a single answer, then stop".

After SFT, the same weights see the prompt wrapped in the template and reliably produce an answer turn followed by <|im_end|>. Nothing about the architecture changes. The model has simply been shown thousands of examples of the behaviour you want.

How does the chat template work?

This series uses a ChatML-style template, which the tokenizer tutorial saved alongside the tokenizer. Every message becomes one turn:

<|im_start|>system
You are a helpful assistant.<|im_end|>
<|im_start|>user
Best poha in Indore?<|im_end|>
<|im_start|>assistant
Try the stalls at Sarafa or Chappan Dukan.<|im_end|>

At inference you send everything up to and including <|im_start|>assistant plus a newline, and the model writes the rest. The template used in training and the one used at inference must match character for character, including newlines.

Why mask the loss to assistant tokens?

If you compute the loss on every token, the model also learns to write user questions and system prompts. Worse, prompts are often longer than answers, so most of the training signal goes into predicting text you'll never ask it to generate.

The fix is a label mask. You keep the full conversation as input, but set the label of every non-assistant token to -100, the value PyTorch's cross_entropy ignores via ignore_index. The <|im_end|> that closes each assistant turn is trained, because that is how the model learns to stop.

The model reads every token but is graded only on its own lines.

Here is the mask built step by step on a tokenized conversation:

tokens = ["<|im_start|>", "system", "\n", "Be", " brief", ".", "<|im_end|>", "\n",
          "<|im_start|>", "user", "\n", "Best", " poha", " in", " Indore", "?", "<|im_end|>", "\n",
          "<|im_start|>", "assistant", "\n", "Try", " Sarafa", ".", "<|im_end|>", "\n"]
ids = list(range(100, 100 + len(tokens)))          # pretend token IDs

labels, in_reply = [], False
for i, tok in enumerate(tokens):
    if in_reply:
        labels.append(ids[i])
        if tok == "<|im_end|>":
            in_reply = False
    else:
        labels.append(-100)                         # ignored by cross-entropy
    if tok == "\n" and tokens[i - 1] == "assistant" and tokens[i - 2] == "<|im_start|>":
        in_reply = True

for tok, lab in zip(tokens, labels):
    print(f"{tok!r:16} {'TRAIN' if lab != -100 else '-'}")
print(f"trained on {sum(l != -100 for l in labels)} of {len(tokens)} tokens")

Only 4 of 26 tokens carry a training signal. Now see what happens to the loss itself when you mask, using made-up per-token losses where the model is already good at the answer but bad at guessing the question:

import numpy as np

# per-token losses for a 20-token chat: 14 prompt tokens, 6 assistant tokens
prompt_losses = np.array([5.1, 4.8, 6.0, 5.5, 4.9, 6.2, 5.7, 5.0, 4.4, 6.1, 5.3, 5.8, 4.7, 5.6])
reply_losses = np.array([1.9, 1.2, 0.8, 1.5, 0.9, 0.4])
mask = np.r_[np.zeros(len(prompt_losses)), np.ones(len(reply_losses))]
all_losses = np.r_[prompt_losses, reply_losses]

unmasked = all_losses.mean()
masked = (all_losses * mask).sum() / mask.sum()
print(f"unmasked loss {unmasked:.2f}  (dominated by guessing the user's question)")
print(f"masked loss   {masked:.2f}  (only the answer you care about)")
print(f"share of gradient from prompt tokens without a mask: {prompt_losses.sum()/all_losses.sum():.0%}")

Without the mask, about 90 percent of the gradient would push the model towards predicting user questions.

The SFT dataset in PyTorch

The dataset class below renders each message separately, so it knows exactly which token IDs belong to the assistant. It then checks the result against the tokenizer's own chat template, pads to a fixed length, and returns inputs and labels already shifted by one.

import json
import torch
from torch.utils.data import Dataset
from transformers import AutoTokenizer

class SFTDataset(Dataset):
    def __init__(self, path, tok, max_len=768):
        self.rows = [json.loads(l)["messages"] for l in open(path, encoding="utf-8")]
        self.tok, self.max_len = tok, max_len
        self.pad = tok.convert_tokens_to_ids("<|endoftext|>")

    def encode(self, messages):
        ids, labels = [], []
        for m in messages:
            head = self.tok.encode(f"<|im_start|>{m['role']}\n", add_special_tokens=False)
            body = self.tok.encode(f"{m['content']}<|im_end|>", add_special_tokens=False)
            tail = self.tok.encode("\n", add_special_tokens=False)
            ids += head + body + tail
            train = m["role"] == "assistant"
            labels += [-100] * len(head) + (body if train else [-100] * len(body)) + [-100] * len(tail)
        return ids, labels

    def __len__(self):
        return len(self.rows)

    def __getitem__(self, i):
        ids, labels = self.encode(self.rows[i])
        ids, labels = ids[:self.max_len + 1], labels[:self.max_len + 1]
        n_pad = self.max_len + 1 - len(ids)
        ids += [self.pad] * n_pad
        labels += [-100] * n_pad
        x = torch.tensor(ids[:-1])
        y = torch.tensor(labels[1:])          # label at t is the token at t+1
        return x, y

tok = AutoTokenizer.from_pretrained("tokenizer")
ds = SFTDataset("data/sft.jsonl", tok)

# sanity check: piecewise encoding must equal the official chat template
msgs = ds.rows[0]
ref = tok.apply_chat_template(msgs, tokenize=False)
assert ds.encode(msgs)[0] == tok.encode(ref, add_special_tokens=False), "template mismatch"
x, y = ds[0]
print(x.shape, (y != -100).sum().item(), "trained tokens")

The assert is worth keeping. If your template and your encoder ever disagree, for example over a missing newline, you'll find out on the first sample instead of after a day of training. Also drop or flag conversations longer than max_len; a reply cut off before <|im_end|> teaches the model to ramble.

SFT training loop and hyperparameters

The training loop is the pretraining script with four changes: load the pretrained weights, read batches from the DataLoader, use a lower learning rate, and run for epochs instead of a fixed number of steps.

import math
import torch
from torch.utils.data import DataLoader
from model import Config, TinyLM

device = "cuda" if torch.cuda.is_available() else "cpu"
ckpt = torch.load("checkpoints/pretrain_last.pt", map_location=device)
cfg = Config(**{**ckpt["config"], "max_seq_len": 1024})   # RoPE table must cover max_len
model = TinyLM(cfg).to(device)
model.load_state_dict(ckpt["model"])

loader = DataLoader(ds, batch_size=16, shuffle=True, num_workers=2, drop_last=True)
epochs, peak_lr, accum = 2, 5e-5, 2
total = epochs * len(loader) // accum
warmup = max(1, total // 50)
opt = torch.optim.AdamW(model.parameters(), lr=peak_lr, betas=(0.9, 0.95), weight_decay=0.0)

def lr_at(s):
    if s < warmup:
        return peak_lr * (s + 1) / warmup
    return 0.1 * peak_lr + 0.45 * peak_lr * (1 + math.cos(math.pi * (s - warmup) / max(1, total - warmup)))

step = 0
model.train()
for epoch in range(epochs):
    for i, (x, y) in enumerate(loader):
        x, y = x.to(device), y.to(device)
        with torch.autocast(device_type=device, dtype=torch.bfloat16):
            _, loss = model(x, y)                     # ignore_index=-100 skips masked tokens
        (loss / accum).backward()
        if (i + 1) % accum == 0:
            for g in opt.param_groups:
                g["lr"] = lr_at(step)
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            opt.zero_grad(set_to_none=True)
            step += 1
            if step % 100 == 0:
                print(f"epoch {epoch} step {step}/{total} loss {loss.item():.3f}")
    torch.save({"model": model.state_dict(), "config": cfg.__dict__}, f"checkpoints/sft_epoch{epoch}.pt")

Starting values that usually work for a 25M to 100M model:

Treat these as a first run. If the model ignores the format after one epoch, raise the learning rate. If it starts forgetting general knowledge it had after pretraining, lower it, or mix a small share of pretraining text back in.

Talk to your chat model

Generation uses the same template with add_generation_prompt=True, and stops at <|im_end|>:

@torch.no_grad()
def chat(model, tok, messages, max_new=120, temperature=0.7, top_k=40):
    model.eval()
    prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    idx = torch.tensor([tok.encode(prompt, add_special_tokens=False)], device=device)
    end = tok.convert_tokens_to_ids("<|im_end|>")
    prompt_len = idx.shape[1]
    for _ in range(max_new):
        logits, _ = model(idx[:, -model.cfg.max_seq_len:])
        logits = logits[:, -1, :].float() / temperature
        kth = torch.topk(logits, top_k).values[:, -1:]
        logits[logits < kth] = -float("inf")
        nxt = torch.multinomial(torch.softmax(logits, -1), 1)
        if nxt.item() == end:
            break
        idx = torch.cat([idx, nxt], dim=1)
    return tok.decode(idx[0, prompt_len:])

print(chat(model, tok, [{"role": "user", "content": "What is the capital of Madhya Pradesh?"}]))

This loop recomputes attention over the whole sequence for every new token, which is fine for testing. Real serving uses a KV cache; for sampling settings see decoding strategies.

Tool calls, thinking and full fine-tuning vs LoRA

The same recipe teaches extra behaviours if your data contains them. Assistant turns with <tool_call>{"name": ..., "arguments": ...}</tool_call> teach a simple JSON tool-call format, and turns with <think>...</think> before the answer teach a visible reasoning step.

Because those tokens were reserved in the tokenizer, nothing else changes. A 64M model learns the format well, but its tool choices and reasoning stay shallow.

You trained every weight here, which is called full fine-tuning and is cheap at this size. For billion-parameter models, LoRA fine-tuning trains small adapter matrices instead.

Common SFT mistakes

• No loss mask. The model starts writing "User:" lines and never answers cleanly.
• <|im_end|> not trained. The model answers, then keeps going forever. Check the end token's label is not -100.
• Template mismatch at inference. A missing newline after assistant is enough to degrade answers.
• Learning rate copied from pretraining. 6e-4 on a small SFT set wipes out what pretraining taught.
• Low-quality data. A small model copies everything, including wrong facts and odd phrasing. See creating an instruction dataset.
• Expecting new knowledge. SFT mostly teaches format and style; facts the base model never learned stay shaky.

Learning to debug fine-tuning runs is easier with someone reviewing your outputs; the Indore GenAI course run by Debugshala is a classroom option.

Key takeaways

• SFT uses the same next-token loss as pretraining, on conversations rendered by the chat template.
• Mask every non-assistant token with -100; train the assistant's reply and its <|im_end|>.
• Use a much lower learning rate than pretraining, 1 to 3 epochs, and sequences long enough for whole chats.
• Assert that your encoding matches the tokenizer's chat template before training.
• Next: improve answers with knowledge distillation and DPO training, then evaluate the model.