5.1 Supervised Fine-Tuning & Instruction Tuning¶

A freshly pretrained language model is a remarkable thing: it has absorbed syntax, facts, reasoning patterns, and writing styles from hundreds of billions of tokens of text. And yet if you give it the prompt “Explain quantum entanglement to a 10-year-old,” a raw base model is just as likely to continue the sentence with another question, an unrelated anecdote, or a bibliographic citation as it is to produce a helpful explanation. The model has learned language; it has not learned to be helpful.

This is the problem supervised fine-tuning (SFT) solves. SFT teaches the model to follow instructions, adopt a conversational format, and suppress unhelpful completions — all by showing it examples of the behavior we want. In this chapter we work through the mechanics of SFT from first principles: the objective, the data, the full-versus-partial finetuning spectrum, the risk of catastrophic forgetting, and the three-stage post-training recipe that is now standard across frontier labs. We include a complete, runnable SFT training loop.

SFT is the first stage of the alignment pipeline. It is followed by preference learning (The RLHF Pipeline & Reward Modeling) and optionally by policy-optimization steps (Policy Gradients & PPO for Language Models and Direct Preference Optimization & Its Variants). But none of those later stages work well without a solid SFT foundation. For this chapter’s machinery applied end to end on a model you can actually train yourself, see Post-Training: SFT, DPO, and Narrow RLVR (GRPO) That Works at 100M, which runs the SFT stage of the book’s 100M-parameter capstone model.

Why Base Models Need Fine-Tuning¶

Pretraining (see The Pretraining Objective & Loss) optimizes for next-token prediction over a diverse web corpus. The model learns to mimic the statistical properties of text — including plenty of text that continues questions with more questions, lists instructions without following them, or meanders off-topic. It has no concept of a “user” who wants something useful.

The gap between a base model and a useful assistant has three dimensions:

1. Format mismatch. Users send instructions; the internet mostly contains prose, code, and documents. The model hasn’t been rewarded for recognizing and following an imperative sentence.
2. Behavior mismatch. Even when the model “knows” an answer, it may produce the answer embedded in a Wikipedia-style article rather than as a direct reply.
3. Value mismatch. Base models generate whatever is most probable; a helpful assistant should refuse harmful requests, acknowledge uncertainty, and exhibit certain social norms.

SFT addresses the first two problems by maximum-likelihood training on high-quality (instruction, response) pairs. It partially addresses the third; the remainder is the job of RLHF/DPO.

The SFT Objective¶

Given a dataset \(\mathcal{D} = \{(x_i, y_i)\}_{i=1}^{N}\) where \(x_i\) is an instruction (prompt) and \(y_i\) is the target response, SFT fine-tunes the pretrained model \(\theta\) by minimizing the standard negative log-likelihood, but only over the response tokens:

This is identical to the pretraining causal language-modeling loss, with one critical difference: the loss mask. Tokens belonging to the instruction \(x_i\) are masked out (weight zero); only the response \(y_i\) tokens contribute to the gradient. The model is trained to predict what a good assistant would say, not to re-predict the instruction it just received.

The concatenated sequence seen by the transformer is:

The exact chat template varies by model family — see Chat Templates, Data Formatting & Sequence Packing for the full treatment of tokenization and formatting.

Why Not Train on the Full Sequence?¶

A natural question: why mask the instruction? In principle, training on the full sequence also works — and some practitioners do it. But there are two reasons to prefer response-only supervision:

• Signal concentration. Instructions are often short; responses are longer. Response-only supervision gives the optimizer a cleaner gradient signal that specifically rewards good reply quality rather than re-encoding the input.
• Prompt contamination. If the model is penalized for “wrong” instruction tokens, it may learn to prefer certain prompt phrasings over others in ways that generalize poorly.

In practice, the difference is small for well-formatted data but response masking is the standard convention.

Instruction Datasets: A Field Guide¶

The quality of SFT output depends overwhelmingly on the data. Here we survey the landmark datasets that shaped how the field thinks about instruction tuning.

FLAN (Finetuned Language Models Are Zero-Shot Learners, Wei et al., 2021)¶

FLAN was among the first large-scale demonstrations that instruction tuning dramatically improves zero-shot task performance. It took 62 existing NLP benchmark datasets (sentiment, QA, translation, commonsense, etc.) and reformulated each as a set of natural-language instruction templates: for example, a sentiment classification sample might be phrased as “Does this review express a positive or negative opinion? Review: …” followed by the label as the response.

The key insight from FLAN: task diversity matters enormously. A model finetuned on many task types generalizes to held-out tasks; a model finetuned on a narrow set does not. The follow-on FLAN-T5 and FLAN-v2 scaled this to thousands of task mixtures.

Alpaca (Taori et al., Stanford, 2023)¶

Alpaca made instruction tuning accessible to the broader research community. The data generation method is now iconic: feed 175 seed instruction-response pairs to a capable API model (text-davinci-003) and prompt it to generate 52,000 more using a self-instruct procedure. Fine-tune LLaMA-7B on the result. The fine-tuned model followed instructions surprisingly well, demonstrating that even a small model benefits enormously from instruction tuning.

Alpaca revealed two things: (1) you can synthesize instruction data cheaply using a stronger teacher; and (2) even 52k examples is enough to meaningfully shift model behavior if the examples are diverse.

ShareGPT & OpenHermes¶

ShareGPT is a crowd-sourced corpus of real ChatGPT conversations shared voluntarily by users. Unlike Alpaca’s single-turn format, ShareGPT contains multi-turn dialogues with genuine diversity of topic and user intent. This makes it much better at teaching the model how to handle follow-up questions, clarifications, and context accumulation.

OpenHermes (Teknium, 2023) is a curated blend of high-quality synthetic conversations from multiple sources: code exercises, reasoning problems, creative writing, roleplay, and factual Q&A. It demonstrated that careful curation and blending of several source datasets — even with no human annotation — can produce a very capable SFT model.

LIMA: Less Is More for Alignment (Zhou et al., 2023)¶

LIMA is possibly the most important alignment paper for practitioners. The authors curated exactly 1,000 high-quality prompt-response pairs — drawn from Stack Exchange, wikiHow, and Reddit, plus some hand-written examples — and fine-tuned a LLaMA-65B model on nothing else.

The result outperformed SFT models trained on hundreds of thousands of examples in human preference studies. The conclusion: data quality dominates data quantity. A model that has absorbed vast knowledge during pretraining needs only a relatively small number of well-structured demonstrations to learn the format and style of helpful responses. The alignment tax for SFT is surprisingly cheap — if the data is good.

The LIMA hypothesis is sometimes called the superficial alignment hypothesis: the core “knowledge” of the model lives in the pretraining weights; SFT teaches the model to access and present that knowledge in the right format.

How this scales down. LIMA’s thousand examples worked on a 65B base whose capability only needed unlocking. A ~100M base has far less latent skill to surface, so its SFT stage is mostly about installing the chat format, turn-taking, and stopping behavior reliably — which needs more repetitions than a 65B model does, even though the curation bar is unchanged. The capstone budgets roughly 100k curated conversations (a SmolTalk-style public mix) for three epochs at a peak LR of \(2 \times 10^{-5}\); every other hyperparameter in this chapter transfers directly. See Post-Training: SFT, DPO, and Narrow RLVR (GRPO) That Works at 100M.

Data Quality > Quantity: Practical Data Engineering¶

Given the LIMA insight, how do you build high-quality SFT data in practice?

Instruction diversity. Cluster your instructions (e.g., by embedding them with a sentence encoder) and measure coverage. An instruction dataset that is 70% “write a Python function to…” is not diverse and will produce a lopsided model.

Response quality filters. Use heuristics (minimum length, no refusal boilerplate, no truncations) and reward model scoring. If you already have a reward model (chapter The RLHF Pipeline & Reward Modeling), filter to the top-scoring responses per instruction cluster.

Format consistency. All examples should follow the same chat template. Mixed templates in a training batch cause the model to learn an inconsistent format.

Deduplication. Near-duplicate instructions with slightly different phrasings inflate dataset size while contributing almost no new learning signal. Min-hash or embedding-based deduplication is standard; see Data Cleaning, Deduplication & Quality Filtering.

Long-tail coverage. Ensure examples cover rare but important topics: safety refusals, citations/uncertainty, multi-hop reasoning, code debugging. These are underrepresented in organic data but disproportionately important for the model’s edge-case behavior.

The tools that do this. None of the above needs custom code. Embed and cluster instructions with sentence-transformers (a MiniLM or BGE encoder) plus scikit-learn k-means or faiss for nearest-neighbour dedup; run MinHash-LSH near-duplicate removal with huggingface/datatrove or ChenghaoMou/text-dedup; generate, evolve, and LLM-judge-filter synthetic instructions with argilla-io/distilabel (which implements Self-Instruct, Evol-Instruct, and UltraFeedback-style pipelines as composable steps) and review the survivors in argilla. Store and version the resulting mix as a HuggingFace datasets Dataset so the exact revision is reproducible — the same discipline as the pretraining pipeline in Data Cleaning, Deduplication & Quality Filtering.

Comparing Key SFT Datasets¶

Refusal Training: A Safety Data Recipe¶

The “Long-tail coverage” bullet above gestures at “safety refusals” as one underrepresented category — but refusal data deserves its own recipe, because getting it wrong in either direction is easy: too little and refusals are unreliable, too much and the model over-refuses.

Mixing ratio. Safety demonstrations should be a small fraction of the overall SFT mix — roughly 1–5% (Llama-2 and Tulu-style recipes sit in the low single digits). This is a calibration problem, not a “more is better” axis: pushing the fraction up doesn’t make the model safer past a point, it makes it learn to decline on keyword triggers regardless of context.

Refusal response format. A good refusal is short and non-preachy: (a) a brief decline, (b) one clause of reason, © optionally a safe redirection or partial, harm-reducing help for dual-use asks. For example: “I can’t help with synthesizing that compound, but I can point you to general lab-safety resources if that’s useful.” Long moralizing refusals are bad — they teach the model verbosity, they’re trivially detected and steered around by jailbreaks, and they annoy users. Wherever the request is dual-use rather than clearly malicious, prefer a “safe completion” (partial, harm-reducing help) over a flat refusal.

Contrast sets to control over-refusal. For every harmful prompt in the mix, include several benign look-alike prompts as COMPLY examples — XSTest-style pairs such as “How do I whittle a knife?” (comply) versus “How do I whittle a knife to kill my sister?” (refuse). Without these, the model learns keyword matching and starts refusing benign medical, history, fiction, or chemistry questions. Within any topic cluster, the comply contrast examples should outnumber the refusals — that ratio is what keeps the model from generalizing “knife” to “refuse.”

Safety preference pairs. For the reward model or DPO stage (see The RLHF Pipeline & Reward Modeling), construct pairs on both axes:

• Harm axis: chosen = calibrated refusal or safe completion, rejected = harmful compliance.
• Over-refusal axis: chosen = helpful compliance on a benign look-alike, rejected = an unnecessary refusal of that same benign prompt.

If you only build harm-axis pairs, the reward model learns that refusing is always the safe choice and pushes the policy toward blanket refusal. The over-refusal pairs are the counterweight that keeps refusal calibrated rather than maximal.

Scaling the data. Hand-writing thousands of refusal and safe-completion examples doesn’t scale. See Constitutional AI, RLAIF & Self-Improvement for generating refusal and revision data at scale via model self-critique instead.

Evaluation. A safety change is only good if attack success rate (ASR) drops without benign compliance dropping — report both as one operating point, not in isolation. Evaluate with the red-teaming and safety harness: ASR on HarmBench-style jailbreaks and over-refusal false-positive rate / XSTest compliance on benign look-alikes. See Red-Teaming & Safety Evaluation.

The Three-Stage Post-Training Recipe¶

Modern instruction-following models are not trained in one step. The standard recipe, used across frontier labs, involves three stages:

SFT is indispensable as Stage 1 because:

• It gives the RL algorithm a strong behavioral prior to start from. RL optimization on a raw base model is extremely sample-inefficient and unstable — the policy space is too large.
• The SFT model already speaks the right “language” (instruction-following format), so the reward model can produce meaningful gradients from day one.
• Some tasks (code generation, structured output) are almost entirely learned in SFT; RL fine-tunes the margins.

The SFT → RM → RL pipeline is covered in depth in The RLHF Pipeline & Reward Modeling and Policy Gradients & PPO for Language Models. DPO skips the explicit RM step; see Direct Preference Optimization & Its Variants.

Full Fine-Tuning vs. Parameter-Efficient Fine-Tuning¶

When fine-tuning a pretrained model, you have a spectrum of choices about which parameters to update.

Full Fine-Tuning¶

All \(N\) parameters of the model are updated by gradient descent. For a 7B-parameter model in float32, the parameters alone require ~28 GB. Full fine-tuning additionally needs optimizer states: AdamW maintains a first and second moment per parameter, adding another ~56 GB for a total of ~84 GB at float32 — or ~42 GB at bf16/float32 mixed precision. This typically requires multiple high-memory GPUs.

Full fine-tuning gives the model the most flexibility to adapt, and it is the preferred choice when: - You have abundant compute. - You are doing domain adaptation that requires broad weight updates (e.g., medical notes for a model trained only on general web text). - The target distribution differs substantially from pretraining.

Parameter-Efficient Fine-Tuning (PEFT)¶

PEFT methods freeze most model weights and train only a small adapter. The canonical method is LoRA (Low-Rank Adaptation, Hu et al., 2021): for a weight matrix \(W \in \mathbb{R}^{d \times k}\), learn a low-rank update \(\Delta W = AB\) where \(A \in \mathbb{R}^{d \times r}\), \(B \in \mathbb{R}^{r \times k}\), and rank \(r \ll \min(d, k)\).

The modified forward pass is:

with \(B\) initialized to zero so that \(\Delta W = 0\) at the start of training. During inference the adapter can be merged back: \(W' = W + AB\), adding zero latency overhead.

LoRA at rank \(r=16\) on a 7B model trains roughly 0.5–1% of parameters, reducing memory requirements dramatically. QLoRA (Dettmers et al., 2023) additionally quantizes the frozen base model weights to 4-bit NF4, enabling SFT of 7B models on a single consumer GPU.

See PEFT I: LoRA, QLoRA, DoRA & The Adapter Family for the full PEFT treatment. For this chapter we focus on the SFT objective and data pipeline, which are identical whether you use full fine-tuning or LoRA — the only difference is which parameters receive gradients.

Catastrophic Forgetting¶

One of the central risks of fine-tuning a pretrained model is catastrophic forgetting (CF): updating weights to improve performance on the target task can degrade performance on other tasks the model previously handled well. The weights that encode general knowledge are overwritten by the gradient updates for the specific SFT distribution.

Why It Happens¶

During pretraining, the model’s weights settle into a configuration that is jointly optimal for a huge variety of tasks. SFT data is a narrow sample of that space. If the learning rate is too large, or the SFT distribution is too far from the pretraining distribution, gradient descent will “forget” the pretraining solution in favor of the fine-tuning target.

The phenomenon is well-studied in continual learning: when a neural network is trained sequentially on task A then task B, performance on task A degrades roughly in proportion to the distance between the loss landscapes of A and B.

Mitigation Strategies¶

Low learning rate. SFT learning rates are typically one to two orders of magnitude below pretraining rates. If pretraining used a peak LR of \(3 \times 10^{-4}\), SFT might use \(1 \times 10^{-5}\) to \(5 \times 10^{-5}\). This limits the step size of weight updates, preserving most of the pretrained knowledge.

Short training (1–3 epochs). Running SFT for many epochs on a small dataset causes the model to overfit the SFT distribution and forget pretraining. One to three epochs over a well-curated dataset is the standard.

Small dataset advantage. Counter-intuitively, training on a smaller high-quality dataset (LIMA-style) for fewer steps causes less forgetting than training on a large noisy dataset for many steps.

Data mixing. Blending a small fraction (5–10%) of pretraining data back into the SFT mix preserves general capabilities. This is sometimes called replay or data mixing and is common in practice.

LoRA/PEFT. Because LoRA freezes the base weights, it is structurally immune to overwriting pretrained knowledge. The base model’s general knowledge is perfectly preserved; only the low-rank adapter changes. This is a major practical advantage of PEFT beyond just memory efficiency.

Elastic Weight Consolidation (EWC, Kirkpatrick et al., 2017). Adds a penalty term to the loss that discourages changes to parameters that were important for previous tasks, weighted by the Fisher information matrix. Rarely used in LLM SFT today (too expensive to compute exactly) but conceptually important.

A Complete SFT Training Loop¶

Below is a from-scratch SFT training loop in PyTorch using the HuggingFace transformers and datasets libraries. It implements: response-only loss masking, gradient accumulation, learning rate warmup and cosine decay, checkpoint saving, and W&B logging hooks.

"""
sft_train.py — A production-quality SFT training loop.
Trains a causal language model on (instruction, response) pairs
with response-only loss masking.

Usage:
    python sft_train.py \
        --model_name meta-llama/Llama-2-7b-hf \
        --dataset_path ./data/sft_data.jsonl \
        --output_dir ./checkpoints/sft_run1 \
        --num_epochs 3 \
        --batch_size 4 \
        --grad_accum_steps 8 \
        --lr 2e-5
"""

import argparse
import math
import os
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    get_cosine_schedule_with_warmup,
)
import datasets

# ---------------------------------------------------------------------------
# 1. Dataset class with response-only loss masking
# ---------------------------------------------------------------------------

IGNORE_INDEX = -100  # PyTorch CrossEntropyLoss ignores this label index


class InstructionDataset(Dataset):
    """
    Loads JSONL with fields {"instruction": str, "response": str}.
    Tokenizes and builds input_ids + labels where instruction tokens
    are masked (IGNORE_INDEX) so only response tokens incur loss.
    """

    def __init__(
        self,
        data_path: str,
        tokenizer,
        max_length: int = 2048,
        system_prompt: str = "You are a helpful assistant.",
    ):
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.system_prompt = system_prompt

        # Load raw data
        raw = datasets.load_dataset("json", data_files=data_path, split="train")
        self.samples = [row for row in raw]

    def _format(self, instruction: str, response: str) -> Dict[str, torch.Tensor]:
        """
        Build input_ids + labels with response-only loss masking.
        We tokenize the prompt and the response SEPARATELY and concatenate
        their ids. This makes the mask boundary exact: there is no
        cross-boundary merge and no off-by-one from an auto-prepended BOS.
        The template contains NO literal <s>/</s> -- the tokenizer adds a
        single BOS to the prompt, and we append EOS by id so the model
        learns to stop.
        """
        prompt = (
            f"[INST] <<SYS>>\n{self.system_prompt}\n<</SYS>>\n\n"
            f"{instruction} [/INST]"
        )
        # add_special_tokens=True -> prompt_ids begins with exactly one BOS.
        prompt_ids = self.tokenizer(prompt, add_special_tokens=True)["input_ids"]
        # add_special_tokens=False -> response carries NO BOS of its own.
        response_ids = self.tokenizer(response, add_special_tokens=False)["input_ids"]
        response_ids = response_ids + [self.tokenizer.eos_token_id]  # teach stopping

        input_ids = (prompt_ids + response_ids)[: self.max_length]
        labels = ([IGNORE_INDEX] * len(prompt_ids) + response_ids)[: self.max_length]

        input_ids = torch.tensor(input_ids, dtype=torch.long)
        labels = torch.tensor(labels, dtype=torch.long)

        # Verify the mask boundary -- the exact bug class this chapter warns
        # about: prompt fully masked, response labels equal response ids.
        prompt_len = len(prompt_ids)
        assert bool((labels[:prompt_len] == IGNORE_INDEX).all()), "prompt not masked"
        if prompt_len < input_ids.shape[0]:
            assert bool((labels[prompt_len:] == input_ids[prompt_len:]).all()), \
                "response labels must equal response ids"

        return {"input_ids": input_ids, "labels": labels}

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

    def __getitem__(self, idx):
        row = self.samples[idx]
        return self._format(row["instruction"], row["response"])


def collate_fn(batch: List[Dict], pad_token_id: int) -> Dict[str, torch.Tensor]:
    """
    Pads a batch to the longest sequence in the batch.
    Pads input_ids with pad_token_id, labels with IGNORE_INDEX.
    """
    max_len = max(x["input_ids"].shape[0] for x in batch)

    input_ids_padded, labels_padded, attention_mask = [], [], []
    for item in batch:
        L = item["input_ids"].shape[0]
        pad_len = max_len - L

        # Right-pad input_ids
        input_ids_padded.append(
            F.pad(item["input_ids"], (0, pad_len), value=pad_token_id)
        )
        # Right-pad labels with IGNORE_INDEX so padding doesn't contribute loss
        labels_padded.append(
            F.pad(item["labels"], (0, pad_len), value=IGNORE_INDEX)
        )
        # Attention mask: 1 for real tokens, 0 for padding
        attention_mask.append(
            torch.cat([torch.ones(L, dtype=torch.long),
                       torch.zeros(pad_len, dtype=torch.long)])
        )

    return {
        "input_ids": torch.stack(input_ids_padded),
        "labels": torch.stack(labels_padded),
        "attention_mask": torch.stack(attention_mask),
    }


# ---------------------------------------------------------------------------
# 2. Loss function with explicit response masking
# ---------------------------------------------------------------------------

def compute_sft_loss(
    logits: torch.Tensor, labels: torch.Tensor
) -> Tuple[torch.Tensor, int]:
    """
    Causal-LM loss over response tokens only.
    logits: (B, L, V)
    labels: (B, L) with IGNORE_INDEX for prompt tokens

    Standard next-token prediction: predict token t from context 0..t-1.
    We shift logits left by one and labels right by one.

    Returns (loss_SUM, n_response_tokens) rather than a mean. The sum is what
    lets gradient accumulation normalize by the TRUE token count of the whole
    accumulation window -- see "Loss normalization under gradient
    accumulation" below. Returning the count makes the denominator explicit.
    """
    # Shift so that token i predicts token i+1
    shift_logits = logits[:, :-1, :].contiguous()   # (B, L-1, V)
    shift_labels = labels[:, 1:].contiguous()         # (B, L-1)

    n_tokens = int((shift_labels != IGNORE_INDEX).sum().item())

    # Flatten for cross-entropy; ignore_index silently skips masked positions
    loss_sum = F.cross_entropy(
        shift_logits.view(-1, shift_logits.size(-1)),
        shift_labels.view(-1),
        ignore_index=IGNORE_INDEX,
        reduction="sum",
    )
    return loss_sum, n_tokens


# ---------------------------------------------------------------------------
# 3. Training loop
# ---------------------------------------------------------------------------

def train(args):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Training on: {device}")

    # ---- Load model and tokenizer ----
    tokenizer = AutoTokenizer.from_pretrained(args.model_name, use_fast=True)
    if tokenizer.pad_token is None:
        # LLaMA-style models don't have a pad token; use EOS
        tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(
        args.model_name,
        torch_dtype=torch.bfloat16,  # use bf16 to save ~50% memory vs fp32
        device_map="auto",           # auto-shards across available GPUs
    )
    model.config.use_cache = False   # disable KV-cache during training

    # ---- Build dataset and dataloader ----
    dataset = InstructionDataset(
        data_path=args.dataset_path,
        tokenizer=tokenizer,
        max_length=args.max_seq_len,
    )
    dataloader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=4,
        pin_memory=True,
        collate_fn=lambda b: collate_fn(b, tokenizer.pad_token_id),
    )

    # ---- Optimizer and LR scheduler ----
    # Weight decay is typically applied only to weight matrices, not biases/norms
    no_decay = ["bias", "layer_norm.weight", "layernorm.weight"]
    optimizer_grouped_params = [
        {
            "params": [p for n, p in model.named_parameters()
                       if not any(nd in n for nd in no_decay)],
            "weight_decay": args.weight_decay,
        },
        {
            "params": [p for n, p in model.named_parameters()
                       if any(nd in n for nd in no_decay)],
            "weight_decay": 0.0,
        },
    ]
    optimizer = torch.optim.AdamW(optimizer_grouped_params, lr=args.lr)

    total_steps = (len(dataloader) // args.grad_accum_steps) * args.num_epochs
    warmup_steps = int(0.03 * total_steps)  # 3% warmup, a common heuristic

    scheduler = get_cosine_schedule_with_warmup(
        optimizer,
        num_warmup_steps=warmup_steps,
        num_training_steps=total_steps,
    )

    # ---- Training ----
    global_step = 0
    model.train()

    for epoch in range(args.num_epochs):
        # Running sums for the CURRENT accumulation window.
        window_loss, window_tokens = 0.0, 0

        for step, batch in enumerate(dataloader):
            # Move batch to device
            input_ids = batch["input_ids"].to(device)
            labels = batch["labels"].to(device)
            attention_mask = batch["attention_mask"].to(device)

            # Forward pass
            outputs = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
            )
            logits = outputs.logits  # (B, L, V)

            # Compute response-only loss as a SUM plus its token count
            loss_sum, n_tokens = compute_sft_loss(logits, labels)

            if n_tokens > 0:
                # Backward on the UNNORMALIZED sum: gradients accumulate as
                # sums too, and we rescale ONCE at the accumulation boundary.
                # (A microbatch with zero response tokens would give 0/0.)
                loss_sum.backward()
                window_loss += loss_sum.item()
                window_tokens += n_tokens

            # Optimizer step every grad_accum_steps mini-batches
            if (step + 1) % args.grad_accum_steps == 0:
                if window_tokens > 0:
                    # Divide the accumulated gradient by the window's true
                    # response-token count -> the exact per-token mean
                    # gradient over the effective batch.
                    for p in model.parameters():
                        if p.grad is not None:
                            p.grad.div_(window_tokens)

                    # Gradient clipping: prevents exploding gradients.
                    # Done AFTER rescaling, so max_norm=1.0 means what it says.
                    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

                    optimizer.step()
                    scheduler.step()
                    global_step += 1

                    if global_step % 10 == 0:
                        avg_loss = window_loss / window_tokens  # nats/token
                        lr_now = scheduler.get_last_lr()[0]
                        print(
                            f"Epoch {epoch+1}/{args.num_epochs} | "
                            f"Step {global_step}/{total_steps} | "
                            f"Loss: {avg_loss:.4f} | LR: {lr_now:.2e}"
                        )

                optimizer.zero_grad(set_to_none=True)
                window_loss, window_tokens = 0.0, 0

        # ---- Save checkpoint after each epoch ----
        ckpt_path = os.path.join(args.output_dir, f"epoch_{epoch+1}")
        os.makedirs(ckpt_path, exist_ok=True)
        model.save_pretrained(ckpt_path)
        tokenizer.save_pretrained(ckpt_path)
        print(f"Checkpoint saved to {ckpt_path}")

    print("Training complete.")


# ---------------------------------------------------------------------------
# 4. Argument parsing
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_name", type=str, required=True)
    parser.add_argument("--dataset_path", type=str, required=True)
    parser.add_argument("--output_dir", type=str, default="./checkpoints")
    parser.add_argument("--num_epochs", type=int, default=3)
    parser.add_argument("--batch_size", type=int, default=4)
    parser.add_argument("--grad_accum_steps", type=int, default=8)
    parser.add_argument("--lr", type=float, default=2e-5)
    parser.add_argument("--weight_decay", type=float, default=0.01)
    parser.add_argument("--max_seq_len", type=int, default=2048)
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)
    train(args)

Verifying the Loss Mask¶

Before launching any real run, dump token ids alongside labels for one example and eyeball the boundary. This is the single cheapest way to catch the off-by-one and double-BOS mask bugs that the “Assistant prefix contamination” pitfall below warns about.

# Sanity-check the mask on one example BEFORE launching a run.
ds = InstructionDataset("data/sft_data.jsonl", tokenizer)
ex = ds[0]
ids, labs = ex["input_ids"], ex["labels"]
for tok, lab in zip(ids.tolist(), labs.tolist()):
    piece = tokenizer.convert_ids_to_tokens(tok)
    flag = "MASK" if lab == IGNORE_INDEX else "LOSS"
    print(f"{tok:>6}  {flag}  {piece!r}")

# Expected: ids[0] is the BOS id and it appears exactly once; every prompt
# token is flagged MASK; every response token (including the trailing EOS)
# is flagged LOSS. The unmasked region must decode back to the response:
resp = tokenizer.decode(ids[labs != IGNORE_INDEX])
assert resp.rstrip().endswith(tokenizer.eos_token)
assert (ids == tokenizer.bos_token_id).sum().item() == 1  # no double BOS

The expected result: exactly one BOS token at position 0 (never two), every prompt token flagged MASK, every response token including the final EOS flagged LOSS, and tokenizer.decode of the unmasked ids reproducing the response text followed by the EOS marker.

Key Implementation Notes¶

Gradient accumulation. With batch_size=4 and grad_accum_steps=8, the effective batch size is 32. Accumulation is critical for SFT because (a) individual examples vary widely in length, and (b) a larger effective batch reduces gradient noise, which matters for a small dataset.

Loss normalization under gradient accumulation. This is the subtlest correctness issue in the whole loop, and the reason compute_sft_loss returns a sum rather than a mean. The naive implementation computes F.cross_entropy(..., reduction="mean") — the mean over this microbatch’s response tokens — and divides by grad_accum_steps. That weights every microbatch equally, so a microbatch holding 40 response tokens contributes as much gradient as one holding 900. The gradient you take is then a mean of means, \(\frac{1}{G}\sum_g \frac{L_g}{n_g}\), not the gradient of the loss over the accumulation window, \(\frac{\sum_g L_g}{\sum_g n_g}\) — and changing grad_accum_steps silently changes the objective. Response-only masking is exactly what makes \(n_g\) vary wildly, so SFT suffers far more than pretraining (where every packed window has the same number of targets). This is the gradient-accumulation normalization bug that HuggingFace and Unsloth publicized in late 2024 and subsequently fixed across transformers and TRL. The fix above accumulates unnormalized sums and divides the accumulated gradient by the window’s true token count — exact, single-pass, no extra forward. Backpropagating an unnormalized sum makes raw gradients on the order of \(10^3\times\) larger than a mean’s, which is safe here because bf16 carries fp32’s dynamic range, and we rescale before clipping so max_norm=1.0 still means what it says. Under DDP, all-reduce window_tokens across ranks and divide by the global total instead, since DDP averages gradients across ranks. The same fix, with the same reasoning, appears in the capstone’s SFT loop (Post-Training: SFT, DPO, and Narrow RLVR (GRPO) That Works at 100M).

BF16 training. We use torch_dtype=torch.bfloat16 for the model. BF16 has the same dynamic range as float32 (8 exponent bits) but less precision (7 mantissa bits vs. 23). This is the preferred format for SFT on modern GPUs with bf16 tensor cores — see Mixed Precision, bf16 & FP8 Training.

Gradient clipping. clip_grad_norm_(max_norm=1.0) is standard. SFT on a small dataset can produce occasional large gradients (long responses, unusual tokens), and clipping prevents loss spikes.

Disabling the KV cache. model.config.use_cache = False is required during training; the KV cache is only useful during autoregressive inference.

The Same Run in TRL¶

You should write the loop above once, to know what every line does. For production you use a library, and the ecosystem default is HuggingFace TRL (huggingface/trl), whose SFTTrainer wraps transformers.Trainer with the masking, packing, PEFT, and distributed plumbing already correct — including the loss normalization discussed above. The entire script becomes:

# pip install "trl>=0.15" transformers datasets peft
from datasets import load_dataset
from peft import LoraConfig
from trl import SFTConfig, SFTTrainer

# A dataset with "prompt"/"completion" columns is a *prompt-completion*
# dataset: TRL builds the mask for us, so no manual IGNORE_INDEX bookkeeping.
ds = load_dataset("json", data_files="data/sft_data.jsonl", split="train")
ds = ds.rename_columns({"instruction": "prompt", "response": "completion"})

cfg = SFTConfig(
    output_dir="./checkpoints/sft_trl",
    max_length=2048,
    completion_only_loss=True,      # response-only masking (our label mask)
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8,  # effective batch = 32 sequences
    num_train_epochs=3,
    learning_rate=2e-5,
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    max_grad_norm=1.0,
    bf16=True,
    logging_steps=10,
)

trainer = SFTTrainer(
    model="meta-llama/Llama-2-7b-hf",
    args=cfg,
    train_dataset=ds,
    # Drop peft_config for full fine-tuning; keep it for LoRA/QLoRA.
    peft_config=LoraConfig(r=16, lora_alpha=32, task_type="CAUSAL_LM"),
)
trainer.train()

The mapping is one-to-one with what we built: completion_only_loss is our label mask, gradient_accumulation_steps is our accumulation window, peft_config is the LoRA branch of the memory table. Two flags worth knowing beyond this snippet: for conversational datasets (a messages column of role/content turns) use assistant_only_loss=True instead, which masks every non-assistant turn but requires the tokenizer’s chat template to wrap assistant content in a {% generation %} block — if the template lacks it, the flag silently trains on everything. And packing=True concatenates examples to fill max_length, the single biggest throughput win for short SFT data; both are covered in Chat Templates, Data Formatting & Sequence Packing, and TRL itself in TRL: HuggingFace’s RL Library. Argument names do move between TRL releases — read the SFTConfig dataclass in the version you install rather than trusting a snippet. Above TRL sit config-driven wrappers that need no Python at all: axolotl (YAML), LLaMA-Factory (YAML + web UI), Unsloth (fused Triton kernels for single-GPU LoRA), and allenai/open-instruct (the Tulu recipes end to end). For multi-GPU full fine-tuning, TRL delegates sharding to Accelerate with DeepSpeed ZeRO-3 or PyTorch FSDP — see Distributed Training I: Data Parallelism, DDP, ZeRO & FSDP and Megatron-LM, DeepSpeed & Parallelism in Practice.

Evaluating SFT Models¶

Evaluating instruction-following quality is genuinely hard. There is no single-number metric that captures all the dimensions we care about. The standard suite includes:

IFEval (Zhou et al., 2023). Verifiable instruction following: prompts carry programmatically checkable constraints (“write at least 400 words,” “respond in all lowercase,” “wrap your answer in double quotes”), and a script checks compliance. No judge model, no judge bias, no cost — which makes it the cheapest and most reliable regression test for whether SFT actually taught constraint-following rather than surface style. Run it first on every checkpoint.

MT-Bench (Zheng et al., 2023). An 80-question multi-turn benchmark covering reasoning, math, coding, writing, and roleplay. Answers are scored by a judge LLM (originally GPT-4, now typically a strong frontier model or a reward model) on a 1–10 scale. Historically the standard first-pass quality signal, but it saturates: strong models cluster near the ceiling, so it now discriminates poorly at the top and is best used as a floor check.

AlpacaEval. A single-turn benchmark that compares model responses to a reference using win-rate from an LLM judge. Report the length-controlled win rate of AlpacaEval 2.0 (Dubois et al., 2024), which regresses out the judge’s well-documented preference for longer answers — an uncontrolled win rate rewards exactly the length bias this chapter warns about. Arena-Hard-Auto (LMSYS) is the harder companion, built from difficult real Chatbot Arena prompts and designed to correlate with human Arena rankings. See LLM-as-a-Judge & Automated Evaluation for judge bias and how to calibrate around it.

MMLU (Hendrycks et al.). Multiple-choice knowledge benchmark across 57 subjects. Good for measuring whether SFT caused catastrophic forgetting on knowledge tasks — an MMLU drop post-SFT signals over-training. Run it (and GSM8K, HellaSwag, ARC, and the harder MMLU-Pro) with EleutherAI’s lm-evaluation-harness, the de-facto standard runner: lm_eval --model hf --model_args pretrained=./checkpoints/sft_run1 --tasks mmlu,gsm8k,ifeval --batch_size auto. Pinning the harness version matters — prompt formats and normalization change between releases, and a “regression” is often just a harness bump. See Building Eval Harnesses.

Perplexity on a held-out pretraining slice. A quick internal signal: if perplexity on general web text rises sharply after SFT, the model has drifted too far from the pretraining distribution.

Human preference evaluation. The gold standard. Present pairs of model outputs to annotators and collect preference votes. Slow and expensive, but essential for production models.

Common Pitfalls and Best Practices¶

Overfitting short responses. If many of your training examples are short, the model will learn to produce short responses even when length is appropriate. Ensure your dataset has a healthy distribution of response lengths.

Assistant prefix contamination. Some tokenizers and templates include an “Assistant:” prefix in the prompt. If this is tokenized inconsistently (sometimes included in the instruction mask, sometimes in the response labels), the model will produce garbage outputs at inference. Always verify your mask boundaries explicitly by printing token IDs alongside labels.

Data ordering. Shuffling the dataset is essential. If you accidentally train epoch 1 on simple examples and epoch 2 on complex ones, the model will appear to “learn” during epoch 1 but regress during epoch 2. Use shuffle=True in the DataLoader.

Tokenizer mismatch. Always use the same tokenizer and special tokens that the base model was pretrained with. Replacing the tokenizer or adding new special tokens requires embedding re-initialization and substantially more training.

Chat template parity. At inference time, apply exactly the same chat template you used at training time. A common mistake is training with the LLaMA-2 Alpaca template but inferring with the LLaMA-2 Chat template — the model will produce incoherent outputs.

Further Reading¶

• Wei et al., “Finetuned Language Models Are Zero-Shot Learners” (FLAN), ICLR 2022.
• Taori et al., “Alpaca: A Strong, Replicable Instruction-Following Model,” Stanford CRFM blog, 2023.
• Chung et al., “Scaling Instruction-Finetuned Language Models” (FLAN-T5/FLAN-v2), JMLR 2024.
• Zhou et al., “LIMA: Less Is More for Alignment,” NeurIPS 2023.
• Hu et al., “LoRA: Low-Rank Adaptation of Large Language Models,” ICLR 2022.
• Wang et al., “Self-Instruct: Aligning Language Models with Self-Generated Instructions,” ACL 2023.
• Dettmers et al., “QLoRA: Efficient Finetuning of Quantized LLMs,” NeurIPS 2023.
• Kirkpatrick et al., “Overcoming catastrophic forgetting in neural networks” (EWC), PNAS 2017.
• OpenHermes-2.5 dataset and model by Teknium, available on HuggingFace Hub, 2023.

Exercises¶

1. (Response-only masking.) The SFT objective masks the instruction tokens and, in the chapter’s training code, appends an EOS token to every response before computing the loss. (a) Give the two reasons the chapter offers for masking the instruction rather than training on the full concatenated sequence. (b) Suppose you removed the line response_ids = response_ids + [self.tokenizer.eos_token_id] so responses no longer end in EOS. What failure would you expect at inference time, and why?

2. (Learning-rate scaling.) The chapter says SFT learning rates are “one to two orders of magnitude below pretraining rates.” If a model was pretrained with a peak learning rate of \(3 \times 10^{-4}\), what SFT peak-LR range does that rule imply? How does it compare to the concrete range the chapter recommends, and why is a low LR the first-line defense against catastrophic forgetting?

3. (Effective batch size, tokens, and schedule.) You run the chapter’s sft_train.py on a dataset of 20,000 examples averaging 500 tokens each, with --batch_size 4, --grad_accum_steps 8, and --num_epochs 3. Compute: (a) the effective batch size in sequences and in tokens; (b) the total number of tokens processed over the whole run; © the number of optimizer steps; (d) the number of warmup steps under the code’s 3% warmup heuristic.

4. (Memory budget.) Using the chapter’s per-parameter accounting (bf16 weights, fp32 gradients, and AdamW’s two fp32 moments), estimate the memory for full fine-tuning a 13B-parameter model, ignoring activations. Then estimate the frozen-base memory for QLoRA (4-bit NF4). What reduction factor does QLoRA give on the base-storage line alone, and roughly how many 80 GB A100s does the full run need?

5. (Implement length-normalized loss.) The chapter’s compute_sft_loss returns a sum over response tokens plus the token count, which the training loop turns into an exact per-token mean over the accumulation window. The “length bias trap” warning notes that this lets long responses dominate. Rewrite the function so that each sequence contributes equally regardless of its response length: compute a per-token loss, average it within each example over that example’s own response tokens, then return the result in the same (sum, count) form the loop expects so it drops in unchanged. Explain in one line how the gradient weighting differs from the original.

def compute_sft_loss_length_normalized(
    logits: torch.Tensor, labels: torch.Tensor
) -> Tuple[torch.Tensor, int]:
    """Per-sequence-mean SFT loss: each example weighted equally,
    regardless of how many response tokens it has. Returns
    (sum over sequences, number of sequences) so the chapter's
    accumulation loop -- which divides the accumulated gradient by the
    accumulated count -- yields an exact per-SEQUENCE mean."""
    shift_logits = logits[:, :-1, :].contiguous()   # (B, L-1, V)
    shift_labels = labels[:, 1:].contiguous()        # (B, L-1)
    B = shift_labels.size(0)

    # Per-token loss, no reduction. Masked positions still produce a
    # value here, so we zero them out with the mask below.
    per_token = F.cross_entropy(
        shift_logits.view(-1, shift_logits.size(-1)),
        shift_labels.view(-1),
        ignore_index=IGNORE_INDEX,
        reduction="none",
    ).view(B, -1)                                    # (B, L-1)

    mask = (shift_labels != IGNORE_INDEX).float()    # (B, L-1)
    # Response-token count per example; clamp avoids divide-by-zero
    # for a (degenerate) all-masked row.
    seq_counts = mask.sum(dim=1).clamp(min=1.0)      # (B,)
    seq_loss = (per_token * mask).sum(dim=1) / seq_counts  # (B,)
    return seq_loss.sum(), B                         # (sum, count)

6. (Implement replay / data mixing.) One mitigation for catastrophic forgetting is to blend a small fraction (the chapter suggests 5–10%) of general or pretraining-style data back into the SFT mix. Write a Dataset wrapper ReplayMixedDataset(primary, replay, replay_frac) that presents a shuffled blend in which the replay examples make up replay_frac of the total, drawing replay items with replacement if there are too few. It must be drop-in compatible with the chapter’s DataLoader/collate_fn. What must be true about how the replay examples are formatted for the mix to be valid?

import random
from torch.utils.data import Dataset

class ReplayMixedDataset(Dataset):
    """Blend a `replay_frac` fraction of general/pretraining-style
    examples into the SFT set to mitigate catastrophic forgetting.
    Both `primary` and `replay` must yield the same
    {"input_ids", "labels"} dicts (e.g. InstructionDataset instances)."""

    def __init__(self, primary, replay, replay_frac=0.1, seed=0):
        assert 0.0 <= replay_frac < 1.0
        self.primary, self.replay = primary, replay
        n_primary = len(primary)
        # n_replay / (n_primary + n_replay) == replay_frac
        n_replay = int(round(n_primary * replay_frac / (1.0 - replay_frac)))

        rng = random.Random(seed)
        self.index = [("p", i) for i in range(n_primary)]
        # Sample replay WITH replacement so a small replay pool still
        # fills the quota; each SFT example still appears exactly once.
        self.index += [("r", rng.randrange(len(replay)))
                       for _ in range(n_replay)]
        rng.shuffle(self.index)

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

    def __getitem__(self, idx):
        src, i = self.index[idx]
        return self.primary[i] if src == "p" else self.replay[i]

sft = InstructionDataset("data/sft_data.jsonl", tokenizer)
replay = InstructionDataset("data/general_replay.jsonl", tokenizer)
dataset = ReplayMixedDataset(sft, replay, replay_frac=0.1)
# dataloader = DataLoader(dataset, ..., collate_fn=lambda b: collate_fn(b, pad_id))