Capstone Lesson 39: Instruction Tuning by Supervised Fine-Tuning

Supervised fine-tuning (SFT) trains a pretrained model on paired examples of an instruction and a desired response, so that it answers instead of continuing. A base model trained on next-token prediction has no idea what an instruction is. Show it the string "What is the capital of France?" and it will continue the question or invent a new sentence. It has the language but not the format contract.

The SFT contract is a string template. Each training example becomes a single sequence with three regions:

<INST> What is the capital of France? <RESP> The capital of France is Paris.

The boundary tokens are special tokens reserved at training time. The model learns that everything after <RESP> is the response, and the response is what gets graded. The base model's next-token objective still applies. It is just trained on a corpus where each example has this shape.

There is a catch. Feed the entire sequence to a vanilla cross-entropy loss and you are training the model to predict the instruction tokens too. The instruction is given, so you want zero gradient on those positions. The fix is the mask.

ignore_index is an argument to torch.nn.functional.cross_entropy. Any target position equal to it contributes zero loss and zero gradient, and PyTorch's convention for the value is -100. The collate function builds two tensors per example: input_ids (the full sequence) and labels (a copy of input_ids with the instruction positions overwritten by -100).

The model sees the whole sequence during the forward pass, so attention can reach the instruction. The loss counts only response tokens. That is the behaviour you want: condition on the instruction, predict the response.

main.py generates two hundred instruction-response pairs deterministically. They cover six task types:

• factual single-shot (capital of X)
• arithmetic
• list extraction
• one-sentence summary
• code (print, sort)
• definition

Each task has a templated instruction and a deterministic response. That is deliberately simple. Exact-match is brittle, so the fixture is built so the right answer is one specific string. Real SFT datasets need fuzzier metrics, but the principle is identical.

The split is 160 train and 40 test. The test set covers all six task types, which lets you report exact-match per category.

The tokeniser is byte-level with three reserved specials:

• INST_ID = 256: marks the start of the instruction region.
• RESP_ID = 257: marks the boundary between instruction and response.
• PAD_ID = 258: padding for variable-length batches.

A sequence looks like [INST] inst_bytes [RESP] resp_bytes [PAD]*. The collate function:

1. Tokenises each example.
2. Pads each example in the batch to the longest sequence in that batch.
3. Builds labels as an unshifted copy of the padded input_ids, with:
   • The instruction region (every position before the response starts) replaced by -100.
   • The padding region replaced by -100.
   • The RESP_ID boundary position itself replaced by -100 (you do not train the model to predict the boundary token; it predicts what follows).

The shift is the standard causal trick, and it happens later, inside shifted_loss. Position i of the logits predicts position i+1, so the loss compares logits[:, :-1] with labels[:, 1:]. The collate function masks the unshifted labels by position, and the slice in shifted_loss then lines each masked label up with the prediction that would have targeted it.

The loop is a standard PyTorch SFT loop: Adam, a learning rate around 3e-4 to 1e-3, ten to twenty epochs on this fixture, no scheduler. The model is small (hidden 96, 2 blocks, max length 96), so it trains to convergence on CPU inside two minutes.

Every fifth epoch (5, 10, 15, 20) the loop runs a quick eval on the held-out set and prints exact-match. run_demo also prints the untrained baseline first, which is 0.0. Watching exact-match climb from that baseline to something like 0.85 by epoch fifteen is the payoff of the lesson. You see the model pick up the format and the answers at the same time.

At eval time the model gets the instruction prefix [INST] inst_bytes [RESP] and generates tokens until either:

• the sequence reaches max_len, or max_new_tokens (64) tokens have been generated,
• the model emits PAD, INST or RESP, or
• the model emits a stop heuristic: two consecutive sentence-ending bytes (., !, ?).

The code ships greedy decoding plus an optional temperature sampler. Exact-match uses greedy, because temperature would make the metric stochastic. Real systems often sample and then judge fuzzily. That pipeline is lesson 41.

Exact-match is the strictest text metric. The predicted response is normalised (lowercase, strip whitespace, collapse double spaces) and compared to the reference response, normalised the same way. Each example scores 1 or 0, and the aggregate is the mean.

Real SFT pipelines pair exact-match with token-level F1 (lesson 41) and a judge model. It stays useful because it is unambiguous: if it says 0.7, exactly 70 percent of test instructions produced the gold response character for character.

The implementation is one main.py plus tests.

1. InstructionTokenizer: byte-level encoder with reserved specials. Encodes either an instruction prefix or a full pair.
2. make_dataset: generates 200 pairs across six task types with a fixed seed.
3. SFTDataset: returns (ids, resp_start) per example. The labels are built later, in sft_collate.
4. sft_collate: dynamic padding, builds the batch tensors, and builds the labels with -100 on instruction, boundary and pad positions.
5. TinyGPT: transformer body plus an untied nn.Linear LM head.
6. train_sft: the SFT loop, with an exact-match eval every eval_every (5) epochs.
7. generate: causal decode from a prefix, greedy or sampled, with the stop rules.
8. exact_match: normalised string comparison, returns float in [0, 1].
9. run_demo: builds the data, trains for twenty epochs, evaluates, prints a per-category breakdown, exits zero on success.

Without the mask, the loss treats instruction tokens as targets, so the model learns to predict the instruction. That is a different objective, and it produces a worse model in two ways. First, capacity goes to reconstructing inputs the user always provides. Second, on a fixture like this one the instruction is often at least as long as the response, so the instruction tokens can take a large share of the loss and gradient. shifted_loss averages over unmasked tokens only, so with the mask the average is taken over response tokens alone and the signal you care about is not diluted. The mask is not polish. It is the objective.

• Add a learning-rate warmup followed by cosine decay. SFT is more sensitive to LR than pretraining.
• Add per-token loss logging and plot the loss curve over training. Early epochs are dominated by template tokens (<RESP>, common prefixes), and later epochs by the actual answer tokens.
• Extend the eval to BLEU-1 or chrF. Exact-match underestimates models that produce a paraphrase with the same answer.
• Add a chat template with multi-turn formatting and train on a fixture that includes follow-ups.

The implementation gives you the format contract, the mask, and the loop. Going from base model to instruction follower comes down to one collate function.