cs336 / lessons/16 · sftlesson 17 / 20

Part V — Alignment & inference

SFT — from completer to assistant

Lesson 14 gave you an honest number for how good a base model is — perplexity down, benchmarks measurable — and lesson 15 gave it long reach. But a well-measured base model with a huge context window is still the wrong kind of program: it maximizes p(text), so it completes, it doesn't answer. This lesson closes that gap with the smallest possible change to the stack — not a new architecture, not a new loss, just a new mask over the loss you already have — and turns a text-completer into something that follows instructions.

The previous step left this broken
You now have a base model that is well-measured (lesson 14) and has long reach (lesson 15): perplexity is down, the benchmarks are honest, and it can attend across tens of thousands of tokens. It clearly has the knowledge and the context to use it. But prompt it with "What is the capital of France?" and it just continues with more questions — the most-likely continuation in its training distribution is often not "Paris" but "What is the capital of Germany? What is the longest river in Europe?", because question-lists are a common shape on the web. The model holds the fact (it can recite "the capital of France is Paris" inside the right context) and now has the reach to condition on a long request, but it has the wrong behavior: it continues text, it does not respond to a request. It has knowledge and reach but not behavior. Evaluation can now see this gap and long context only widens what it could answer; neither can fix it. Nothing in pretraining ever told the model that a turn ends and the other party should now speak.
Linear position
Forced by: a base model is a text-completer, not an instruction-follower — it has the knowledge but not the behavior, and no amount of better measurement changes that.
New idea: supervised fine-tuning (SFT) — keep the exact pretraining cross-entropy loss, but train on (prompt, response) pairs wrapped in a chat template and mask the loss to the response tokens only. Data quality ≫ quantity (LIMA: ~1,000 strong examples). The behavior is "superficial" in the precise sense that it teaches format and style, not new facts.
Forces next: SFT can only make the model imitate a demonstrated answer; it has no way to say one answer is better than another, to learn from cheap comparisons, or to exceed its demonstrators — which forces preference optimization (lesson 17).
The plan
Six moves. (1) Pin down the behavior gap and the superficial-alignment hypothesis. (2) Show the loss is unchanged from lesson 2 — only the mask is new — and where every gradient term goes. (3) Build the chat template as a tokens-only protocol: roles, multi-turn, example packing, and the classic bug. (4) Get the data right: quality, diversity, and how little is enough. (5) Make it cheap: parameter-efficient fine-tuning (LoRA / QLoRA), which is how SFT is actually done in practice. (6) Drive the mask yourself, then hit the wall SFT cannot climb, which hands off to preferences.

1 · The behavior gap, and why it is shallow

A base model is a near-optimal next-token predictor for its corpus. That makes it a brilliant completer and a useless assistant, and the two are different objects. Completion answers "what text is most likely to follow this prefix?" Assistance answers "what is the helpful response to this request?" On the open web those targets disagree constantly: a question is far more often followed by another question (an FAQ, a quiz, a forum thread) than by its answer, so the most-probable continuation of "What is the capital of France?" is structurally a worse answer than the one the model could give if it were asked to.

The fix you might reach for first — "train it on more facts" — is the wrong diagnosis. The knowledge is already in there: pretraining on terabytes of text has seen "Paris is the capital of France" thousands of times, and the model will happily produce it in a context that looks like a context where that sentence belongs (an encyclopedia paragraph, say). What is missing is the protocol: a learned convention that "a user turn just ended; now produce the helpful answer and then stop." That convention is what SFT installs.

The superficial-alignment hypothesis (Zhou et al., 2023, "LIMA")
A model's knowledge and capabilities are learned almost entirely during pretraining; alignment mostly teaches it which sub-distribution of formats and styles to surface when interacting with a user. The strong evidence: LIMA fine-tuned a 65B base model on just ~1,000 carefully curated prompt/response pairs (no RLHF) and got an assistant competitive with far more heavily tuned models. If alignment were installing new capability, 1,000 examples could not possibly do it; the fact that it can means SFT is unlocking a behavior the weights already supported. Treat this as a working hypothesis, not a law — it is strongest for style and instruction-following and weakest for genuinely new skills (e.g. a format the base model never saw), which need more data.

The practical consequence sets up everything below: if SFT is mostly format and style, then a small set of high-quality, on-distribution demonstrations should beat a large pile of mediocre ones — and that is exactly what the data section finds.

2 · The loss is unchanged — only the mask is new

Recall lesson 2: the model minimizes the next-token cross-entropy summed over every position of a sequence,

Lpretrain = − Σt log pθ(xt+1 | x≤t)

SFT minimizes the same quantity over the same model, with one difference: the sum runs only over positions inside the response. Concatenate each pair into one sequence s = [prompt tokens, response tokens] and let mt ∈ {0,1} be 1 exactly on response positions:

LSFT = − Σt mt · log pθ(xt+1 | x≤t)

That is the entire algorithmic change. No new architecture, no new optimizer, no new vocabulary, no new sampling protocol, no second model. The same AdamW + cosine schedule (lesson 4), the same FlashAttention forward pass (lesson 7), the same FSDP sharding (lesson 8) — everything from Parts I–II carries over verbatim. You change two things only: the data (now formatted as prompt/response turns) and which positions count toward the loss (only the response). For the line-by-line, off-by-one view of the mask tensor in real code, see gpt_mini / 03 · SFT.

Why prompt tokens still run forward but contribute no gradient
The prompt positions are not skipped in the forward pass — they cannot be, because every response token attends back over the whole prompt (that is how it conditions on the request). What the mask removes is their contribution to the scalar loss: a masked position adds 0 to L, so it adds 0 to the gradient. The model reads the prompt; it is never asked to predict it.

Here is the masking in one slice of PyTorch-ish code. IGNORE_INDEX = −100 is the default value F.cross_entropy drops from the loss:

s      = prompt_ids + response_ids          # one concatenated sequence, length T = P + R
x      = s[:-1]                             # inputs   (length T-1)
y      = s[1:].clone()                      # targets  (length T-1), shifted by one
y[:P-1] = IGNORE_INDEX                      # mask every target that lies inside the prompt
loss   = F.cross_entropy(logits.view(-1, V), y.view(-1), ignore_index=IGNORE_INDEX)

The shift-by-one is the subtle part: the first response target sits at index P−1 of y (because y[i] is the prediction made from input x[i] = s[i]). So y[:P-1] masks exactly the prompt-internal predictions and keeps the all-important "prompt just ended, now begin the answer" transition in the loss. Writing y[:P] would silently mask that first response token — the single most behavior-defining prediction — and is one of the most common SFT bugs.

3 · The chat template — a tokens-only protocol

The model has exactly one input channel: token ids. So "this span is the user's request; that span is your answer; the turn ends here" must all be expressed in tokens. That is the chat template: a small set of reserved special tokens that delimit roles, learned by the model the same way it learns any other token — by repeated exposure during SFT. There is no hardware support and no privileged channel; the roles are a convention written in the byte stream.

A widely used scheme is ChatML: a turn is <|im_start|>role … <|im_end|> with roles system, user, assistant. A single training example looks like this (every line is just tokens; the markers are single special-token ids, not literal angle-bracket characters):

<|im_start|>system
You are a helpful assistant.<|im_end|>
<|im_start|>user
What is the capital of France?<|im_end|>
<|im_start|>assistant
The capital of France is Paris.<|im_end|>
       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  <- loss is masked ON here (and on <|im_end|>)
                                                   everything above is masked OFF

Three roles, three jobs. system sets standing behavior (persona, rules) and is part of the prompt — masked off. user carries the request — masked off. assistant is what the model must learn to produce — masked on, including the closing <|im_end|>, because learning to stop is as important as learning what to say; an assistant that never emits its end token rambles forever at inference.

Multi-turn. A conversation is just more turns concatenated in order. Every assistant turn is masked on; every system/user turn is masked off — so one multi-turn example yields several supervised response spans, each conditioned on the full prior context via attention. A 4-turn dialogue trains all four assistant replies in a single forward pass.

Example packing. SFT examples are short and variable-length; padding each to the context window wastes most of the compute budget (lesson 5's MFU). Instead you pack several examples back-to-back into one T-length sequence. The mask already handles correctness — prompt spans are off, response spans are on, regardless of how many examples share the sequence. The one thing packing must also fix is cross-example attention: a block-diagonal / document mask so tokens of example B cannot attend to example A. Forget that mask and you leak one conversation into the next; the loss still looks fine, so the bug is invisible until you notice degraded multi-turn behavior.

The classic bug: training on the prompt
By far the most common SFT mistake is computing the loss over the whole packed sequence — prompt tokens included — as if it were pretraining. The gradient at a user-token position pushes the model to predict user text ("after a user types 'What is the', the next token should be 'capital'"). You are now partly training a model that generates the user's side of the conversation, which dilutes the response signal and degrades instruction-following. Signal: the SFT loss is suspiciously low and smooth from step 0 (predicting prompts is easy, on-distribution text), and the resulting model tends to echo or continue the user's phrasing instead of answering. The widget below lets you turn this bug on and watch the consequence.

4 · Data — quality ≫ quantity

If alignment is mostly surfacing existing behavior (§1), then the lever is which behavior you demonstrate, not how much. The empirical picture across the field:

Quality dominates
LIMA: ~1,000 hand-curated examples ≳ tens of thousands of noisy ones. A few thousand excellent demonstrations is a strong instruction-tuned model; the marginal mediocre example can actively hurt by teaching sloppy format.
Diversity over volume
Coverage of task types (QA, summarize, rewrite, code, refuse, multi-turn) matters more than count within any one type. The model generalizes the protocol; it needs to see the protocol exercised across the space of requests.
How much is enough
Order ~1k–100k examples for instruction-following — tiny next to pretraining's trillions of tokens. Past a few epochs SFT starts to overfit style and forget pretraining breadth; 1–3 epochs is typical.
Where it comes from
Human-written demonstrations (expensive, best), curated public instruction sets (FLAN, Dolly, OpenAssistant), and distilled outputs from a stronger model (Alpaca-style — cheap, scalable, but caps quality at the teacher and inherits its errors).

Two cautions worth internalizing. First, SFT teaches what it is shown — including formatting tics, hedging, and refusals. If every demonstration opens with "Certainly! Here is…", your model will too. Second, the superficial-alignment hypothesis has a sharp edge: because SFT mostly redistributes existing behavior rather than adding capability, demonstrating an answer the base model cannot actually produce (a fact it never learned, a format it never saw) teaches it to imitate the shape of such answers — i.e. to confidently make things up. SFT on beyond-its-knowledge demonstrations is one mechanism by which fine-tuning can increase hallucination. Demonstrate behavior the weights already support.

5 · Parameter-efficient fine-tuning (LoRA)

So far "fine-tune" has meant full SFT: every one of the model's N parameters gets a gradient and is updated, and you store a complete N-parameter copy of the model per task. At the budget of this course that is brutal — a separate full checkpoint for the coding assistant, the summarizer, the support bot — and most of that storage is re-encoding a base model that barely moved. The superficial-alignment hypothesis (§1) hints at the fix: if alignment only nudges the model into a sub-distribution it already supports, the change SFT makes to the weights should be small and low-dimensional. So why pay to update all N of them?

LoRA (Low-Rank Adaptation) takes that seriously. Freeze the pretrained weight matrix W (no gradient, no optimizer state) and learn only a low-rank update alongside it: W + ΔW, where ΔW = B·A with A ∈ ℝr×d, B ∈ ℝd×r, and the rank r ≪ d (often 8–64 against a model dimension in the thousands). You train only A and B — typically well under 1% of the parameters — which collapses optimizer memory (no Adam moments for the frozen weights) and the per-task storage to a tiny adapter instead of a full model.

Two payoffs follow from the shape of ΔW = B·A. At inference you can merge the product B·A back into W once, giving a single matrix with zero added latency versus the base model — LoRA is not a runtime tax. Or you keep the adapters separate and hot-swap many of them on one frozen base: dozens of task-specific behaviors served from a single set of shared weights, each a few megabytes. This is exactly the multi-tenant serving setup the inference lesson cares about.

QLoRA pushes memory further by stacking two ideas: quantize the frozen base to 4-bit (the precision idea from lesson 18 · inference, also central to the distillation track's model-compression toolkit), then train the small LoRA adapters in higher precision on top of it. The frozen weights cost a quarter of the memory and never need gradients, so the dominant cost vanishes — enough to fine-tune a 65B model on a single GPU, something full SFT of that model could never fit.

The trade-off is honest but usually mild: restricting the update to rank r is strictly less expressive than a full-rank update, so on tasks that genuinely move the weights a lot LoRA can trail full fine-tuning slightly. But SFT is exactly the regime where the change is small and superficial (§1) — which is why parameter-efficient fine-tuning, not full SFT, is how instruction-tuning and alignment are actually done in practice once you leave the from-scratch setting of this course.

6 · Drive the mask yourself

The widget renders one ChatML example as role-colored token chips. The toggles let you choose which spans count toward the loss. The default (assistant span only) is correct SFT. Turn on the prompt or system spans and you are committing the classic bug from §3 — watch the "% tokens in loss" climb and the verdict flip from "instruction-follower" to "text-completer," because most of your gradient is now spent learning to predict the user's text rather than the answer. Turn the assistant span off and there is no supervised signal at all.

Loss-mask visualizer — which tokens does SFT learn from?
Green outline = this token contributes to the loss (gradient flows). Toggle a span to include/exclude it. Correct SFT = only the assistant span is on. Unmasking system/user spans is "training on the prompt": the KPI shows the loss budget draining into prompt tokens, and the verdict degrades. The chips reuse the shared .tok roles — gray = prompt, blue = assistant.
Tokens in loss
% tokens in loss
% of gradient on prompt
Resulting model

Notice the asymmetry the KPIs make concrete: the response is a small fraction of the sequence (here about a third (9 of 29 tokens)), so correct SFT puts all its gradient on that small slice. The moment you unmask the prompt, the majority of every gradient step is spent on tokens you never wanted the model to generate. The mask is not a tuning detail — it is the entire difference between an assistant and a completer.

Failure modes & checklist

Failure modes

  • Training on the prompt. Loss over the whole sequence, not just the response. Signal: SFT loss is low and smooth from step 0, and the model echoes or continues the user's phrasing instead of answering it.
  • Off-by-one mask (y[:P] not y[:P-1]). Masks the first response token — the prompt→answer transition. Signal: the model is fine mid-answer but unreliable at starting to answer; it sometimes drifts back into completion mode at the turn boundary.
  • Forgetting the stop token. Mask off the closing <|im_end|> and the model never learns to end its turn. Signal: coherent answers that then ramble on into a fabricated next turn until the length cap.
  • No cross-example mask when packing. Tokens attend across packed examples. Signal: loss looks normal, but multi-turn / context-following quality is mysteriously worse than the data suggests.
  • SFT beyond the model's knowledge. Demonstrating answers the base model can't actually produce. Signal: increased confident hallucination on exactly the topics you "taught."
  • Template mismatch at inference. Train with one chat template, serve with another. Signal: a model that benchmarks well but behaves erratically in the product because the role tokens it learned aren't the ones it's being fed.

Checklist

  • Mask everything but the response. System and user spans contribute 0 to the loss; assistant span (incl. its end token) is on.
  • Verify the shift. First trained target is at index P−1; draw it on paper once.
  • Pack with a document mask. Block-diagonal attention so packed examples can't see each other.
  • Curate, don't hoard. A few thousand diverse, high-quality, on-distribution demonstrations beats a large noisy set.
  • Stay inside the model's knowledge. Demonstrate behavior the weights already support; 1–3 epochs.
  • Freeze one template. Train and serve with the identical chat template and special tokens.

Checkpoint

Try it
Open the widget with the default (assistant span only) and read off the "% tokens in loss" — it should be a minority of the sequence, and "% of gradient on prompt" should be 0. Now tick include user span: watch the percentage jump and the verdict flip to "text-completer." Predict, before you click, what ticking include system span on top of that does to the gradient-on-prompt figure. Finally, turn the assistant span off entirely and explain in one sentence why the resulting "0 tokens in loss" means the SFT step is a complete no-op — the gradient is exactly zero and θ never moves.

Where this points next

SFT got you an instruction-follower with one mask line. But look at what it can and cannot express. It can say "produce this demonstrated answer" — it imitates. It cannot say "answer A is better than answer B," because cross-entropy on a single demonstration has no notion of a competing answer to rank against. It cannot learn from a comparison (which is far cheaper to collect than a gold demonstration — humans rank more reliably than they write). And it cannot exceed its demonstrators: the very best SFT can do is reproduce the distribution of answers it was shown. To optimize relative quality — to push the model toward better-than-demonstrated answers using a preference or a verifier signal, while a KL anchor keeps it from drifting into gibberish — you need a different objective. That is preference optimization, and the deeper RL treatment lives in reinforcement_learning / 14 · the LLM era. Next: 17 · Preferences & rewards — RLHF, DPO, GRPO.

Takeaway
A base model maximizes p(text), so it completes rather than answers — it has the knowledge but not the behavior. SFT repairs exactly that, with the smallest change in the whole stack: keep the identical pretraining cross-entropy and the identical model, change only the data (prompt/response pairs wrapped in a chat template's role tokens) and the mask (loss on response tokens only, including the stop token). The behavior it installs is "superficial" in LIMA's precise sense — format and style, not new capability — which is why ~1,000 strong examples can do it and why quality and diversity beat raw volume. The classic failure is training on the prompt: a wrong mask spends most of your gradient teaching the model to generate the user's text. SFT's hard ceiling is that it can only imitate a demonstration; it has no way to rank one answer above another, to learn from cheap comparisons, or to surpass its demonstrators — which is the wall that forces preference optimization in lesson 17.

Interview prompts