The Foundations of LLMs, Part II: Supervised Fine-Tuning

by David Oniani

Posted on

Table of Contents

This is the second installment in the series, The Foundations of LLMs. In the previous article, we built and pretrained a small transformer from scratch in PyTorch. Pretraining and midtraining teach a model how language works by having it read large amounts of text and play a simple game: guess the next token. Formally, for a sequence of tokens \((x_1, \dots, x_T)\) of length \(T\), a model parameterized by weights \(\theta\) minimizes the negative log-likelihood of each target token \(x_t\) given all preceding context tokens \(x_{\lt t} = (x_1, \dots, x_{t-1})\):

$$\mathcal{L}_{\text{pretrain}}(\theta) = -\sum_{t=1}^T \log P_{\theta}(x_{t} \mid x_{\lt t})$$

While this game teaches the model grammar, facts, and reasoning, the resulting base model is still just a text completer, not an assistant. When you type a question, the model does not know that you want an answer, so it simply guesses what words would most likely follow your sentence in a book or on a webpage. If you give it a homework question, it might reply with three more homework questions, because questions on the internet frequently appear inside numbered worksheets and exam banks.

The SFT Objective

To turn this text completer into a helpful assistant, we use Supervised Fine-Tuning (SFT), also known as instruction tuning. In SFT, we train the model on curated examples of a user asking a question (the prompt) and an assistant giving a helpful answer (the response). The training objective looks almost identical to pretraining, with one critical change to which token positions are summed:

$$\mathcal{L}_{\text{SFT}}(\theta) = -\sum_{t \in \mathcal{T}_{\text{response}}} \log P_{\theta}(x_{t} \mid x_{\lt t})$$

Here, \(\mathcal{T}_{\text{response}} \subset {1, \dots, T}\) is the subset of target token positions that belong to the assistant's reply (including its closing end-of-turn marker). During SFT, we only grade the model on the assistant's answer tokens and ignore its guesses on the user's question so the model spends all of its capacity learning how to answer rather than imitating noisy user prompts. Unlike the causal attention mask in Part I (which hides future tokens during the forward pass), loss masking happens only at the end when scoring the output: the model still reads and attends to the entire user question via \(x_{\lt t}\), but gradients are computed only on the reply.

Completion vs. Instruction Following. Because a base model is trained to continue arbitrary documents, it treats a question as the start of a worksheet or article rather than a request to fulfill. Conditioning the model on a dedicated assistant turn that is graded strictly on direct answers shifts its behavior from document continuation to instruction following:

  • Answering a Question
    • Input: "Why is the ocean salty?"
    • Pretrained Model: "Why is the sky blue? Why do volcanoes erupt?"
    • SFT Model: "Rainwater erodes rocks and washes dissolved minerals into the ocean."
  • Following an Instruction
    • Input: "Write a Python function to reverse a string."
    • Pretrained Model: "...without using built-in methods. Problem 2: Write a function..."
    • SFT Model: "def reverse(s: str) -> str: return s[::-1]"

Chat Templates and SFT Tensors

During pretraining, the model reads one long, unbroken stream of text. To have a back-and-forth conversation, the model needs clear delimiters showing who is speaking and when each turn ends, a structure called a chat template. Production models with subword tokenizers reserve dedicated special token IDs like "<|im_start|>user", "<|im_start|>assistant", and "<|im_end|>" that the tokenizer refuses to produce from raw user input, preventing a user from spoofing a turn boundary by typing "Assistant:" inside their prompt. In our character-level transformer from Part I, we format each turn using plain text labels made of characters already in our vocabulary: "\nUser: {prompt}\nAssistant: {response}\n". Here, "\nUser: " marks the start of the human query, "\nAssistant: " signals that it is the model's turn to reply, and the closing newline "\n" acts as our end-of-turn marker.

Recall from Part I that Transformer.forward computes training loss using PyTorch's F.cross_entropy via F.cross_entropy(logits.view(B * T, -1), targets.view(-1)). By default, F.cross_entropy sets ignore_index=-100 and reduction="mean": whenever a target position equals -100, PyTorch excludes that position from both the forward sum and the denominator, dividing the total loss only by the number of unmasked response tokens \(\lvert \mathcal{T}_{\text{response}} \rvert\) rather than \(B \times T\). Normalizing by \(\lvert \mathcal{T}_{\text{response}} \rvert\) ensures that a long user prompt never dilutes the gradient magnitude on a short assistant answer. To build a training example, we create an input sequence token_ids and a target sequence targets shifted forward by one character, and set every prompt target to -100:

def create_sft_tensors(
    prompt: str,
    response: str,
    tokenizer: CharacterLevelTokenizer,
    device: torch.device,
) -> tuple:
    """Prepares aligned input and target tensors for Supervised Fine-Tuning."""

    formatted_prompt = f"\nUser: {prompt}\nAssistant: "
    full_text = f"{formatted_prompt}{response}\n"

    prompt_token_ids = tokenizer.encode(formatted_prompt)
    full_token_ids = tokenizer.encode(full_text)

    token_ids = torch.tensor(full_token_ids[:-1], device=device).reshape(1, -1)  # Input sequence
    targets = torch.tensor(full_token_ids[1:], device=device).reshape(1, -1)     # Shifted by +1
    targets[:, : len(prompt_token_ids) - 1] = -100                               # Mask prompt tokens

    return token_ids, targets

Let us verify both boundaries of this tensor slice. First, why do we mask up to len(prompt_token_ids) - 1 instead of len(prompt_token_ids)? Because targets is shifted left by one index (full_token_ids[1:]), targets[0, i] holds the character full_token_ids[i + 1] that token_ids[0, i] must predict. At index i = len(prompt_token_ids) - 1, token_ids[0, i] is the final character of formatted_prompt (the trailing space in "\nAssistant: "), and its target targets[0, i] is full_token_ids[len(prompt_token_ids)], which is the very first character of response and must therefore remain unmasked. Second, at the right boundary, full_token_ids[:-1] drops the closing newline "\n" from token_ids (since nothing is predicted after the turn ends), while full_token_ids[1:] keeps that closing "\n" as the final unmasked target in targets so the model is explicitly graded on learning when to stop talking. This exact same rule extends to multi-turn conversations by masking every user turn with -100 and leaving every assistant reply unmasked.

Why SFT Works with So Little Data

While pretraining processes trillions of tokens across thousands of GPU-hours to teach a model language, world facts, and reasoning from random initialization, SFT can align that pretrained base model into a conversational assistant using only a few thousand curated examples and a small fraction of the compute budget. Two core dynamics explain this sample efficiency:

Preparing the Frankenstein Dataset

Because SFT only formats knowledge the model already has, our training examples must match what our Part I model read during pretraining: Mary Shelley's Frankenstein. If we try to fine-tune that tiny model on Python code or geography trivia, it will fail because those concepts do not exist in its weights, and if a prompt contains any character absent from Frankenstein, CharacterLevelTokenizer.encode will raise a KeyError. In addition, because each formatted question-and-answer turn in our dataset is 89 to 107 characters long, feeding it into Part I's default block_size=32 model would immediately raise an IndexError when indexing self.pos_emb_table(torch.arange(T)) (and we cannot simply resize pos_emb_table after pretraining because rows 32 through 127 would contain untrained random noise). We therefore pretrain Part I's Config with block_size=128 and construct four in-domain question-and-answer pairs using vocabulary drawn directly from Shelley's novel:

frankenstein_sft_data = [
    {
        "prompt": "Who created you?",
        "response": "Victor Frankenstein fashioned me in his laboratory.",
    },
    {
        "prompt": "Where were you brought to life?",
        "response": "In a solitary chamber at Ingolstadt, on a dreary night in November.",
    },
    {
        "prompt": "Why do you seek revenge?",
        "response": "Because mankind cast me out into solitude and misery.",
    },
    {
        "prompt": "What do you demand from your creator?",
        "response": "A companion of the same species to share my exile.",
    },
]

The Training Loop

Because create_sft_tensors already sets the prompt targets to -100, we do not need to change a single line of our Transformer class from Part I. Notice that because each question-and-answer pair has a different sequence length \(T\), we process one conversation at a time (B = 1). Batching multiple variable-length conversations together (B > 1) in production pipelines requires either right-padding shorter sequences (setting padded target positions to -100 and masking padded positions in attention) or sequence packing (concatenating multiple conversations into a single block_size window with a block-diagonal attention mask so one conversation cannot attend to another). Here, running 64 epochs over our 4 examples performs \(64 \times 4 = 256\) single-example gradient steps (roughly 25,000 characters total, over 600 times less compute than Part I's 16.4-million-character pretraining run):

def train_sft(
    model: Transformer,
    dataset: list[dict[str, str]],
    tokenizer: CharacterLevelTokenizer,
    optimizer: torch.optim.Optimizer,
    device: torch.device,
    epochs: int = 64,
) -> None:
    """Fine-tunes the transformer on conversational instruction-response pairs."""

    model.train()
    for epoch in range(1, epochs + 1):
        total_loss = 0.0
        for item in dataset:
            token_ids, targets = create_sft_tensors(item["prompt"], item["response"], tokenizer, device)

            optimizer.zero_grad()
            _, loss = model(token_ids, targets)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        avg_loss = total_loss / len(dataset)
        print(f"Epoch {epoch:2d}/{epochs} | Average Loss: {avg_loss:.4f}")

# Assuming `model` (pretrained with `block_size=128`), `tokenizer`, and `device` are from Part I
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)
train_sft(model, frankenstein_sft_data, tokenizer, optimizer, device)

Conversational Inference

Once fine-tuning is finished, we can chat with our model using the generate method from Part I. To ask a question, we wrap the prompt in our template ending with "\nAssistant: " and leave the reply blank. Seeing "\nAssistant: " cues the model to generate the assistant's answer, and we set max_new_token_ids=96 so even our longest 67-character answer fits with room to spare. Because Part I's minimal generate loop always samples a fixed number of tokens without an early stop-token check (whereas production inference engines pass an eos_token_id to break out of the loop as soon as the end-of-turn token is sampled), the model will keep generating characters after emitting its closing newline "\n", so we slice .split("\n")[0] to return only the assistant's single turn:

def ask_assistant(
    prompt: str,
    model: Transformer,
    tokenizer: CharacterLevelTokenizer,
    device: torch.device,
    max_new_token_ids: int = 96,
) -> str:
    """Formats an evaluation prompt and generates an assistant response."""

    eval_prompt = f"\nUser: {prompt}\nAssistant: "
    token_ids = torch.tensor(tokenizer.encode(eval_prompt), device=device).reshape(1, -1)
    output = model.generate(token_ids, max_new_token_ids=max_new_token_ids, temperature=0.7)[0]
    full_text = tokenizer.decode(output.tolist())
    return full_text[len(eval_prompt) :].split("\n")[0]  # Stop at the closing newline

reply = ask_assistant("Who created you?", model, tokenizer, device)
print(f"Assistant: {reply}")

On our 140K-parameter miniature model trained on 4 examples, querying one of the training prompts recalls the target response ("Assistant: Victor Frankenstein fashioned me in his laboratory."), while querying an unseen prompt produces the learned single-turn assistant format (terminating cleanly at "\n") populated with blended Frankenstein phrases, since a 140K character model lacks the capacity to synthesize novel factual answers. On a frontier base model pretrained on trillions of tokens, by contrast, fine-tuning on a small set of demonstrations generalizes the assistant behavior to arbitrary unseen questions across any domain stored in the base weights.

Conclusion

Supervised Fine-Tuning turns a raw next-token predictor into a helpful conversational assistant. By wrapping text in clear "\nUser: " and "\nAssistant: " role markers and masking the user's question from the loss function with -100, we teach the model to answer questions directly without changing a single line of our Part I transformer architecture. While our four-example Frankenstein dataset is a minimal demonstration, the exact same masked cross-entropy objective powers SFT across frontier post-training pipelines. Still, SFT has a fundamental limitation: because it trains the model via teacher forcing on positive reference answers only, it cannot show the model what not to say, penalize plausible-sounding hallucinations, or score the overall helpfulness and safety of a complete response. To address this, frontier models follow SFT with preference alignment stages such as Reinforcement Learning from Human Feedback (RLHF) (which trains a reward model on human comparisons and optimizes the language model policy, meaning the probability distribution \(P_{\theta}\) choosing each token action, via reinforcement learning) or Direct Preference Optimization (DPO) (which directly increases the likelihood of preferred responses over rejected responses in closed form).