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][tfol-pretraining], we explored the architecture, pretraining, and midtraining of an autoregressive transformer from scratch in PyTorch. Pretraining over web-scale text and midtraining on curated reasoning data produce a capable completion engine. By optimizing the maximum likelihood objective over text corpora, the model learns the statistical distribution of language. Formally, it minimizes the negative log-likelihood of a sequence:

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

While this objective yields deep syntactic understanding and broad semantic knowledge, the resulting base model remains fundamentally unconstrained. When presented with an input prompt, it does not attempt to answer your question or solve your problem. Instead, it merely generates the tokens most likely to naturally follow that string on the internet. If given an exam question, it might output a sequence of additional questions. If given an incomplete thought, it will simply continue generating prose.

The SFT Objective

To transform this open-ended completion engine into a helpful AI assistant, we apply Supervised Fine-Tuning (SFT), also known as instruction tuning. SFT conditions the network on high-quality, structured instruction-response pairs. The objective function looks remarkably similar to pretraining, but with one critical twist:

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

Notice the index set \(\mathcal{T}_{\text{response}}\). In SFT, gradient updates are applied exclusively to the target response tokens. By masking out the loss for the prompt itself, we preserve the network's ability to read and process instructions while aggressively optimizing its ability to generate the correct answers.

Understanding the Gap: Completion vs. Instruction Following. The following concrete examples illustrate the behavioral gap between a raw causal base model and an instruction-tuned model:

  • Answering a Direct Science Question
    • Input: "Why is the ocean salty?"
    • Pretrained Model: "Why is the sky blue? Why do volcanos erupt? Chapter 4 explores..." (the base model treats the prompt as an item in a textbook question bank)
    • SFT Model: "The ocean is salty primarily because rain breaks down continental rocks, washing minerals and dissolved ions into rivers that flow out to sea." (the instruction-tuned model provides a direct factual answer)
  • Fulfilling a Coding Request
    • Input: "Write a Python function to check whether a string is a palindrome."
    • Pretrained Model: "...and then test it using pytest. In our previous article, we looked at arrays..." (the base model treats the prompt as a snippet inside a tutorial article)
    • SFT Model: "def is_palindrome(text: str) -> bool: return text == text[::-1]" (the instruction-tuned model writes the requested code implementation)
  • Adopting an Instructive Persona
    • Input: "Explain neural networks to a fifth grader using a kitchen analogy."
    • Pretrained Model: "Explain neural networks to college students using linear algebra." (the base model outputs parallel prompts from an assignment sheet)
    • SFT Model: "Think of a neural network like a kitchen full of taste testers who adjust a cookie recipe until the cookies bake perfectly every time." (the instruction-tuned model adopts the metaphor)
  • Summarizing Technical Information
    • Input: "Summarize: Photosynthesis converts solar energy into chemical sugars, absorbing carbon dioxide and emitting oxygen as a byproduct."
    • Pretrained Model: "Plants require sunlight, water, and nutrients found in mineral-rich soil..." (the base model treats the text as an opening paragraph and elaborates)
    • SFT Model: "Photosynthesis captures sunlight to create sugars from carbon dioxide, releasing oxygen." (the instruction-tuned model produces a concise summary)

Constructing SFT Tensors and Chat Scaffolding

In base pretraining, text is fed into the model as unstructured character or subword streams. To turn this into a conversation, post-training relies on chat templates, wrapping messages in distinct structural delimiters. In frontier LLMs with subword tokenizers, these delimiters are represented as special tokens like <|im_start|>user and <|im_start|>assistant.

In our character-level transformer from Part I, we implement the identical architectural principle using unambiguous text markers that exist directly in our Frankenstein vocabulary:

\nUser: {prompt}\nAssistant: {response}\n

These structural markers act as an unambiguous conditioning signal, telling the model when a human is speaking and when it is expected to reply.

To train the model, we pack the formatted prompt and target response into aligned input and target tensors. Recall from Part I that our Transformer.forward method accepts an optional targets argument:

# From Part I Transformer.forward:
if targets is None:
    logits = self.proj(out[:, [-1], :])
    loss = None
else:
    logits = self.proj(out)
    loss = F.cross_entropy(logits.view(B * T, -1), targets.view(-1))

In PyTorch, F.cross_entropy includes a built-in ignore_index parameter that defaults to -100. When the loss function encounters a target value of -100, it completely excludes that position from both the forward loss calculation and backward gradient backpropagation.

Because causal language modeling predicts token t + 1 from tokens up to position t, our input tensor consists of all tokens except the very last one, and our target tensor consists of all tokens shifted forward by one position. In the target tensor, every token belonging to the prompt is masked with -100, while the response tokens remain untouched.

Here is how we construct these specialized SFT training tensors to work directly with our CharacterLevelTokenizer and Transformer from Part I:

import torch

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

    # 1. Format the conversation using structural role markers
    formatted_prompt: str = f"\nUser: {prompt}\nAssistant: "
    full_text: str = f"{formatted_prompt}{response}\n"

    # 2. Tokenize the prompt and full conversational turn
    prompt_tokens: list[int] = tokenizer.encode(formatted_prompt)
    full_tokens: list[int] = tokenizer.encode(full_text)

    # 3. Create next-token prediction slices
    # Input sequence: all tokens except the very last one
    input_ids: torch.Tensor = torch.tensor(full_tokens[:-1], dtype=torch.long, device=device).unsqueeze(0)
    # Target sequence: all tokens shifted forward by one position
    targets: torch.Tensor = torch.tensor(full_tokens[1:], dtype=torch.long, device=device).unsqueeze(0)

    # 4. Mask the prompt tokens in the target
    # The first response token is predicted from the prompt at index len(prompt_tokens) - 1
    targets[:, : len(prompt_tokens) - 1] = -100

    return input_ids, targets

Two details in this construction deserve attention:

First, notice the trailing \nAssistant: header at the end of the prompt. During inference, we feed the prompt ending with this prefix and no subsequent text. This acts as a clear turn-taking cue, prompting the transformer to step into the assistant role and begin generating the answer.

Second, notice the trailing newline \n at the end of the full conversation. In our character-level vocabulary, this terminal character functions as the End-Of-Sequence marker. By training the model to predict this closing newline at the end of every response, the network learns exactly when to conclude its answer and yield its turn.

When expanding this concept to multi-turn conversations, the masking logic naturally generalizes. In a multi-turn exchange, all user turns are masked with -100 in the target tensor, while cross-entropy loss is computed across the assistant responses. Alternatively, conversations can be unfolded into separate training records where earlier dialogue turns serve as unmasked conditioning context and the network is trained to predict the latest assistant reply.

Loss Masking Mechanics

With create_sft_tensors producing pre-aligned inputs and masked targets, the training step plugs directly into our Transformer class from Part I without modifying a single line of the model architecture:

# Assuming 'model', 'tokenizer', and 'device' are initialized from Part I
input_ids, targets = create_sft_tensors(
    prompt="Why is the ocean salty?",
    response="Rain erodes rocks and carries dissolved salts through rivers to the sea.",
    tokenizer=tokenizer,
    device=device,
)

# Forward pass: Part I Transformer computes masked cross-entropy internally
_, loss = model(input_ids, targets)

# Backpropagation
loss.backward()

Inside Transformer.forward, providing targets causes the model to project the full sequence into vocabulary logits (self.proj(out)) and execute F.cross_entropy. Because all prompt positions in targets are set to -100, the prompt generates zero loss and zero gradients. The prompt acts purely as causal conditioning context, updating network weights exclusively to match the target response distribution.

Why SFT Works: Alignment and Optimization

A remarkable property of SFT is its striking sample efficiency. Pretraining and midtraining require huge compute budgets and massive corpora to instill broad world knowledge, syntax, and reasoning capabilities. In contrast, effective SFT can often be accomplished with a small fraction of that computational footprint.

Several core dynamics explain this contrast:

Context Window and Vocabulary Considerations

To ensure seamless execution with the Part I architecture, two configuration choices require attention:

  1. Context Window Sizing (block_size): In Part I, we initialized our base transformer with block_size = 32 for rapid demonstration on a single CPU. However, multi-character prompt-response pairs easily exceed 32 tokens. Following the context length extension concepts introduced during midtraining in Part I, our fine-tuning model is configured with an expanded context window (such as block_size = 128), ensuring that the full conversation fits within the positional embedding table and causal attention mask.

  2. In-Domain Vocabulary Compatibility: In Part I, our CharacterLevelTokenizer extracted its vocabulary directly from Project Gutenberg's Frankenstein. If an instruction prompt contains characters outside this vocabulary, the encoder will raise an error. Structuring our role markers as standard text (\nUser: and \nAssistant: ) guarantees that all formatting tokens already exist in the model's vocabulary and embedding weights.

In-Domain SFT Dataset

This brings up a crucial constraint of SFT: it is designed to elicit behavior and format responses, not to inject brand new world knowledge.

In Part I, our miniature model was trained exclusively on Mary Shelley's Frankenstein. If we attempt to fine-tune it on generic instruction datasets (such as Python coding problems or general geography trivia), it will fail because those facts do not exist in its weights. To demonstrate SFT cleanly on our 140K-parameter transformer, we condition the model on concepts it already learned during its Frankenstein pretraining:

frankenstein_sft_data: list[dict[str, str]] = [
    {
        "prompt": "Who created you?",
        "response": "Victor Frankenstein fashioned me in his laboratory.",
    },
    {
        "prompt": "Where were you brought to life?",
        "response": "In a garret at Ingolstadt, amid boiling cauldrons.",
    },
    {
        "prompt": "Why are you so vengeful?",
        "response": "Because mankind abandoned me to loneliness and misery.",
    },
    {
        "prompt": "What do you demand from your creator?",
        "response": "A companion of the same species to share my exile.",
    },
    {
        "prompt": "Who was Victor's closest friend?",
        "response": "Henry Clerval, who accompanied him on his travels.",
    },
]

Every sample in this dataset fits within our block_size = 128 context window and draws strictly from the vocabulary acquired during pretraining.

The Training Loop

With the dataset and tensor preparation ready, we construct the full training loop. We encapsulate the fine-tuning procedure inside train_sft, optimizing parameters using torch.optim.AdamW with a conservative learning rate of 5e-5. Notice how each step calls _, loss = model(input_ids, targets) directly, mirroring the pretraining loop from Part I:

import torch.optim as optim

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

    model.train()
    for epoch in range(epochs):
        total_loss: float = 0.0

        for item in dataset:
            optimizer.zero_grad()

            # 1. Prepare aligned tensors
            input_ids, targets = create_sft_tensors(
                prompt=item["prompt"],
                response=item["response"],
                tokenizer=tokenizer,
                device=device,
            )

            # 2. Forward pass through Part I model
            _, loss = model(input_ids, targets)

            # 3. Backpropagation and optimizer step
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

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

# Assuming 'model', 'tokenizer', and 'device' are already initialized from Part I
optimizer: optim.AdamW = optim.AdamW(model.parameters(), lr=5e-5)
train_sft(
    model=model,
    dataset=frankenstein_sft_data,
    tokenizer=tokenizer,
    optimizer=optimizer,
    device=device,
    epochs=15,
)

Conversational Inference

Once fine-tuning completes, we evaluate our aligned model using the generate method implemented in Part I. At test time, we supply the prompt formatted with our assistant role prefix and allow the model to complete the turn autoregressively:

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

    # Format the prompt with the trailing assistant turn-taking cue
    eval_prompt: str = f"\nUser: {prompt}\nAssistant: "
    input_ids: torch.Tensor = torch.tensor(
        tokenizer.encode(eval_prompt),
        dtype=torch.long,
        device=device,
    ).unsqueeze(0)

    # Autoregressive generation using Part I generate method
    output_ids: torch.Tensor = model.generate(input_ids, max_new_token_ids=max_new_tokens, temperature=0.7)
    full_text: str = tokenizer.decode(output_ids[0].tolist())

    # Extract the assistant response following the role cue
    return full_text[len(eval_prompt) :]

# Query the fine-tuned model
reply: str = ask_assistant(
    prompt="Who created you?",
    model=model,
    tokenizer=tokenizer,
    device=device,
)
print(f"Assistant: {reply}")

Because the model was trained with masked cross-entropy conditioned on the \nAssistant: prefix, it immediately generates text as the creature, demonstrating successful instruction alignment.

Conclusion

SFT transforms an unconstrained autoregressive completion engine into an instruction-following assistant. By pairing chat templates with selective response loss masking, we steer the model's pretrained representations without destroying its underlying knowledge. In this tutorial, we constructed specialized training tensors, integrated them seamlessly with our Part I transformer, and fine-tuned our base model on in-domain Frankenstein trivia.

While SFT instills the desired conversational format, frontier models often undergo further alignment using Reinforcement Learning from Human Feedback (RLHF) or Direct Preference Optimization (DPO). These subsequent stages refine nuanced preferences, safety, and helpfulness, completing the journey from raw pretraining to a fully aligned AI assistant.