Chat templates and loss masking
A team fine-tunes a base model on 50,000 chat transcripts. The training loss falls nicely. Then they talk to it, and two strange things happen. Sometimes the model answers and then keeps going, writing <|im_start|>user and a plausible next question in the user's voice, and answering that too. Other times it never stops at all, rambling until it hits the token limit.
Both bugs live in a few lines of preprocessing that decide which tokens the loss is computed on. In the first case every token was trained, user turns and role headers included, so the model learned to play both sides. In the second the end-of-turn token was masked, so the model never learned to emit it. This section is about those few lines. We'll derive the SFT loss as a masked version of the pretraining loss, work out exactly which positions to keep in single-turn, multi-turn, and tool-using conversations, find a quiet averaging bug that changes which examples matter, pack many short examples into one sequence without letting them see each other, and find the response tokens reliably when a real tokenizer is involved.
From a conversation to one sequence of tokens
Section 2.5 showed how a chat template serializes a list of messages, with role markers, contents, and end-of-turn tokens, into a single token sequence. Here is the ChatML layout again, for a conversation with a system message and two exchanges.
<|im_start|>system\nBe brief.<|im_end|>\n <|im_start|>user\nCapital of France?<|im_end|>\n <|im_start|>assistant\nParis.<|im_end|>\n <|im_start|>user\nAnd Italy?<|im_end|>\n <|im_start|>assistant\nRome.<|im_end|>\n
Every token of that sequence belongs to one of a few kinds. There are header tokens (<|im_start|>, the role name, the newline), content tokens of a particular role, end-of-turn tokens (<|im_end|>), and separators (the newline after <|im_end|>). A loss mask is a decision, per kind and per role, about which of them the model should be trained to produce. Let's build the sequence with a character-level stand-in for a tokenizer, so that every token is visible, and record each token's role and kind.
Of 100 tokens, 13 are trained, the assistant's two answers and their two end-of-turn tokens. Everything else is context the model reads but is never asked to write. Now let's see why that is the right choice, and what it means for the loss.
The SFT loss is the pretraining loss with weights
Write the rendered conversation as one sequence . Pretraining would minimize the average negative log-likelihood of every token given the ones before it (§6.3). For SFT we want something else. Write the model as , the policy notation of §14.1 (it is the same language model, read as a map from a prompt to a distribution over responses). In the single-turn case with prompt and response , the sequence is the concatenation , and the chain rule splits its log-likelihood in two,
In plain words, the pretraining loss on a conversation asks the model to predict the user's words as well as its own answer. SFT keeps only the second term, , because we want a good assistant, not a good simulator of users. In token terms, give each target position (the position whose target is ) a weight , and minimize
Pretraining is everywhere. SFT is when is a trained token and otherwise, so with the set of counted target positions this is exactly the masked cross-entropy of §6.3, , implemented by setting the label of every untrained position to . Values in between are a prompt loss weight, on prompt targets, with as standard SFT and as pretraining.
The gradient is where masking can be misunderstood. Differentiating,
In plain words, a masked position contributes nothing to the gradient as a prediction. But the prompt tokens are still in the forward pass. Every response position attends to them (§4.4), so the gradient of each response term flows back through attention into the prompt positions' hidden states and the weights that computed them. The model is still trained to read the prompt well, just not to write it. Masking the loss and masking attention are different operations, and SFT needs only the first.
Why not train on the prompt too?
Two reasons, one about effort and one about capacity.
Effort. With everywhere and normalization by the total count, the response's share of the gradient is its share of the tokens. In a retrieval-augmented example with 2,000 tokens of retrieved documents and a 100-token answer, under 5% of the gradient goes to the answer, so learning the behavior you care about is about 20 times slower at the same learning rate.
Capacity. A model trained on prompts spends parameters modeling the prompt distribution, the way users write or what documents say, and those parameters are shared with the response. We can see the interference in the smallest possible model, a bigram table (one next-token distribution per previous token), with prompts drawn from one Markov chain (users and documents) and responses from another (the assistant). For a bigram the maximum-likelihood solution is just normalized transition counts, so we can compare the two losses exactly.
With short prompts the difference is modest (2.26 against 2.05 nats). With 2,000-token prompts the fully trained model is far worse at the response, 4.38 nats, because its predictions are now mostly a model of the prompt chain. A Transformer is not a bigram table, and it can learn to condition on role tokens, but the pressure is the same, and the effort argument holds for any model.
The evidence on real models is mixed in an informative way. Llama 2's SFT zeroed the loss on prompt tokens (Touvron et al. 2023), Self-Instruct set OpenAI's prompt loss weight to 0, and QLoRA's ablations found training on the target alone better for MMLU across four datasets (Dettmers et al. 2023). On the other side, Shi et al. (2024) found that keeping the loss on instructions helps when instructions are long and outputs short, or when the dataset is small (LIMA-sized), and attributed the gain to less overfitting, and Huerta-Enochian and Ko (2024) found small nonzero prompt weights (0.01 to 0.5) best for short-completion data. The prompt loss is a regularizer. Default to masking it, and treat a small as a hyperparameter when your responses are short and your data is small.
Aligning the mask with the shift
The next-token shift is where off-by-one bugs are born, so let's be precise. There are two conventions in use, and you must know which one your code follows.
- Pre-shifted. You build inputs
x = ids[:-1]and targetsy = ids[1:], as our GPT did in §6.1, and the target at position is the token at . The mask goes on the targets, so position is counted when token is a trained token. - Unshifted labels. Hugging Face models take
labelsaligned withinput_idsand shift internally (logits[:, :-1]againstlabels[:, 1:]). You setlabels[i] = -100when token itself is not trained.
The two are the same rule expressed at different offsets. Here is the first assistant turn of our conversation, with the label for each input position under the pre-shifted convention.
Look at the newline after assistant. It is a header token, so it is never trained as a target, but it is the input at the position whose target is "P", the first letter of the answer. That position must be counted. If you mask by the role of the input token instead of the target token, you silently drop the first token of every response (the model never learns how to start an answer) and keep the label after the end-of-turn token (the model learns to predict \n after <|im_end|>, harmless) or worse, under other templates, the first token of the next user turn.
Multi-turn conversations
A conversation with assistant turns could be cut into training examples, the first ending after turn 1, the second after turn 2, and so on, each trained only on its last turn. Training all assistant turns of the full conversation in one sequence gives exactly the same per-token losses. The causal mask guarantees that the prediction at any position depends only on earlier positions (the induction argument of §4.4), so the loss on turn 's tokens in the long sequence is the loss you would get from the truncated example ending at turn . One forward pass trains all examples, at the cost of one. That is why "train every assistant turn" is the default.
Sometimes you want less, and the mask is the tool.
- Last turn only. When earlier assistant turns were produced by a different, weaker system (scraped chat logs, an old model), you may want them as context but not as targets. Llama 2's Ghost Attention did this on purpose. It attached a system instruction ("act as Napoleon") to every user turn of a synthetic dialogue, sampled the latest model's responses, then removed the instruction from all but the first turn and set the loss to zero on all earlier turns, assistant messages included, so that only a response conditioned on a long-ago instruction was trained (Touvron et al. 2023).
- Selected turns. Corrective data shows a mistake and its fix. Train on the fix, not on the mistake, by masking the mistaken turn while keeping it visible as context.
- Stripped reasoning. Reasoning models often drop the chain of thought of earlier turns from the history at inference (the gpt-oss model card says that in multi-turn conversations the reasoning traces of past assistant turns should be removed). Training must match. Render earlier turns without their reasoning and train the reasoning only on the final turn, or the model sees histories at training time that it will never see in deployment.
System prompts
A system message sets the rules for the conversation, and the model reads it but never writes it, so its tokens are always masked. What matters is the distribution of system prompts in training. If every training example has the same system prompt, or none, the model learns that system prompts are decoration. To make a model actually follow them, train on many different system prompts paired with responses that obey them, including constraints that persist over many turns, which is what Ghost Attention was built to teach.
Three practical details bite. Templates may inject a default system prompt when you supply none (Qwen2.5's template inserts "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."), so your data and your serving must make the same choice. Some families have no system role at all (Gemma 2's template raises an error), and system content goes into the first user turn. And models deployed behind an API increasingly learn an instruction hierarchy, a ranked order of trust between message sources. The gpt-oss models were trained with the order System, then Developer, then User, then Assistant, then Tool, so that a tool output cannot override the developer (OpenAI 2025). That ordering is learned from data in which lower-priority messages try to override higher ones and the trained response refuses to comply, which ties this section to prompt injection defenses (§19.3).
Tool calls
A tool-using conversation adds two kinds of message, the model's request to call a tool and the tool's result. Here is Qwen2.5's real template applied to a one-tool conversation, rendered from the template shipped with the model.
<|im_start|>system
You are Qwen, created by Alibaba Cloud. You are a helpful assistant.
# Tools
You may call one or more functions to assist with the user query.
You are provided with function signatures within <tools></tools> XML tags:
<tools>
{"type": "function", "function": {"name": "get_weather", "description": "Current weather for a city", "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}}}
</tools>
For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
<tool_call>
{"name": <function-name>, "arguments": <args-json-object>}
</tool_call><|im_end|>
<|im_start|>user
Is it raining in Oslo?<|im_end|>
<|im_start|>assistant
<tool_call>
{"name": "get_weather", "arguments": {"city": "Oslo"}}
</tool_call><|im_end|>
<|im_start|>user
<tool_response>
{"temp_c": 7, "rain": true}
</tool_response><|im_end|>
<|im_start|>assistant
Yes, it is raining in Oslo, and it is 7 degrees.<|im_end|>
The tool definitions are rendered into the system message, the call is an assistant turn containing JSON between <tool_call> tags, and the tool's result comes back as a user turn wrapped in <tool_response> tags. The masking rule follows from who produces each piece. The model must learn to emit the call, so the assistant's tool-call tokens and their end-of-turn are trained. The tool response comes from the environment, so it is masked, always. A model trained on tool outputs learns to write them, which at inference means hallucinating a weather report instead of waiting for the real one. Note that in this template the tool response sits inside a turn whose role header says user, so a mask built from role names alone would treat it like any user turn. That happens to be right here, but a mask should be built from the message structure (which message each token came from) rather than by parsing rendered text. Other families use dedicated roles or tokens (ipython in Llama 3.1, a tool role and channels in gpt-oss), and the rule is the same everywhere. Train what the model emits, mask what the world returns.
Multi-turn loss mask with shifted targets
Implement build_example(messages, encode, last_turn_only=False) for the ChatML layout of this section, with the special ids in SPECIAL and the byte-level encode from the setup. Each message is a dict with "role" and "content", and an optional boolean "train". Render each message as <|im_start|> + encode(role + "\n") + encode(content) + <|im_end|> + encode("\n").
A token is trained if it is the content or the <|im_end|> of an assistant message whose "train" flag is not False (missing means True). With last_turn_only=True, only the last assistant message (among those not flagged False) is trained. Messages with role "tool", "user", or "system" are never trained.
Return (x, y), the pre-shifted inputs and targets as integer NumPy arrays of length , where x = ids[:-1] and y[t] = ids[t+1] if token is trained, else -100.
Per-token or per-example averaging
The normalization in hides a choice. Take a batch of examples, where example has trained tokens with per-token losses . There are two natural averages.
In plain words, the first gives every trained token the same weight, . The second gives every example the same weight, so a token in example gets weight . A token in a 5-token answer counts a hundred times more than a token in a 500-token answer. The two agree only when all are equal.
Neither is wrong, but they are different objectives. The token average is maximum likelihood of the corpus of response tokens. The example average says every conversation matters equally, however long, and so it upweights short answers. What is wrong is getting one when you meant the other, and that happens silently. Gradient accumulation computes a token mean inside each micro-batch, so averaging micro-batch losses produces a hybrid that weights micro-batches equally (§6.3). Data parallelism does the same across devices if each device normalizes by its own token count (§10.2). Tulu 3's authors traced a gap between two of their training setups to exactly this. Their fix was a sum loss, dividing by nothing (or by a constant) and retuning the learning rate, so that every token is weighted equally regardless of how the batch is split (Lambert et al. 2024). Hugging Face's Trainer was fixed in 2024 to pass the total number of counted items across accumulation steps into the loss, and TRL's average_tokens_across_devices exists for the same reason.
The 5-token example gets 0.9% of the gradient under token averaging and a third of it under example averaging. Split the batch into two micro-batches and average their means, and each of the first two examples' tokens gets weight 0.009, the 500-token example's tokens get 0.001, which matches neither objective. Pick one deliberately and make sure your loss computes it under every batch layout.
Masked cross-entropy, by token or by example
Implement sft_loss(logits, targets, reduction) for logits of shape (B, T, V) and targets of shape (B, T) with marking ignored positions. Return (loss, dlogits), the scalar loss and its gradient with respect to the logits.
reduction="token"averages the per-token negative log-likelihood over all counted tokens in the batch.reduction="example"averages within each example over its counted tokens, then averages over the examples that have at least one counted token. Examples with no counted token contribute nothing and are not counted in .
It must be stable for logits near 1000, and ignored positions must get exactly zero gradient.
Packing examples into one sequence
SFT examples vary enormously in length. A one-line answer and a ten-turn coding session can sit in the same batch, and padding every row to the longest wastes most of the compute (§6.1 measured how much, about 70 percent for short documents in long rows). Packing concatenates several examples into one row of length . Three things must then be true for it to be exact, that is, for every example to get the same loss as it would alone.
- Attention stays inside each example. A block-diagonal causal mask allows query to attend to key only if both belong to the same example and .
- Positions restart at 0 for each example, so the position encodings match the unpacked case. With learned absolute positions, as in our small models, this is required, or an example that starts mid-row is read at the wrong positions. With RoPE (§5.5), attention scores depend only on the offset between two tokens, so once the block mask keeps examples apart the reset changes nothing in the math (the point §6.1 made for pretraining). Pipelines reset positions anyway, because the reset is how padding-free kernels find the boundaries.
- No target crosses a boundary. The last token of example A must not be trained to predict the first token of example B. With per-example labels already masked, this means setting the target at the last position of every example to , which in the pre-shifted convention happens automatically if you shift each example before concatenating.
Production systems never build the mask. They pass the cumulative sequence lengths, cu_seqlens = [0, n_1, n_1 + n_2, ...], to a variable-length attention kernel that skips the off-diagonal blocks entirely (FlashAttention's varlen interface, §12.4). Hugging Face added this to its training stack in Transformers 4.44 as a collator that flattens a batch into one row and passes position_ids, from which the boundaries are recovered where positions reset to 0 (Kundu et al. 2024). Llama 2's SFT filled its 4,096-token sequences by concatenating all the prompts and answers of the training set, with a special token separating prompt and answer segments (Touvron et al. 2023), and for short documents many pipelines still accept the cross-contamination. For SFT, where examples are short and many share a row, the exact version is cheap and removes a confound.
Pack SFT examples with block-diagonal attention
Implement pack(examples, T). Each example is a pair (x, y) of pre-shifted inputs and targets (equal-length integer lists, targets may contain ). Place the examples in order into one row of length T (assume they fit, total length at most T), padding the rest with input 0 and target . Return a dict with
"x","y", integer arrays of lengthT,"pos", position ids that restart at 0 for each example (padding positions get 0),"seq", the example index of each position, and-1for padding,"cu_seqlens", the integer array[0, n_1, n_1 + n_2, ..., total](padding not included),"mask", the(T, T)boolean block-diagonal causal mask. Padding positions attend only to themselves, so that no row is empty.
Because the examples are pre-shifted, no target crosses a boundary, and your packing must not create one.
Finding the response tokens in a real tokenizer
Our render knew the role of every token because it built the sequence one message at a time, byte by byte. A real pipeline starts from a template string and a BPE tokenizer, and there the mask is easy to get subtly wrong. There are three ways to find which tokens came from the assistant, and they are not equally safe.
Search for a response marker. Older recipes rendered the conversation, tokenized it, and searched the ids for the tokens of a marker such as ### Assistant: (TRL's former completion-only collator worked this way). It breaks because tokenization does not commute with concatenation (§2.5). The marker tokenized alone is not the marker tokenized after a newline. Here is the failure with a greedy longest-match tokenizer standing in for BPE, whose vocabulary, like real ones, has a token for a newline followed by ###.
The search fails. A collator that finds no marker masks the whole example, so it trains on nothing and nothing crashes (TRL's collator printed a warning and dropped the example from the loss). The patch was to tokenize the marker together with the text that precedes it, which works until someone changes the template.
Tokenize each message and concatenate the ids. This is what our render does. It is exact when every boundary falls on a special token that BPE never merges across, which is one reason chat templates put a special token at every turn boundary. Where a boundary is ordinary text, the concatenated ids can differ from those of the whole string, and the model is trained on token sequences it will never see at inference.
Tokenize once and map character offsets. Render the whole string, record the character span of each assistant message while you render it, tokenize once with offsets (return_offsets_mapping=True in Hugging Face's fast tokenizers), and train every token whose character range overlaps an assistant span. This is exact by construction, because the tokens are the ones inference will see. It is what the last line of the cell does, and it finds Paris, ., and <eos>. Templates that mark assistant content with {% generation %} blocks hand you the spans, and apply_chat_template(..., return_assistant_tokens_mask=True) performs this mapping.
One rule needs a decision. A token can straddle the boundary, starting in the header and ending inside the answer. The model must produce that token to begin its answer, so it is trained. The rule is overlap, at least one shared character, not starts inside the span.
Truncation is part of masking. An example longer than the maximum length gets cut, and cutting from the right removes the end of the response, its end-of-turn token first. If that happens often, the model learns from answers that never end. If the whole response is cut, the example has no trained tokens, and under per-example averaging it is a 0/0. Filter or split over-long examples, or truncate the prompt from the left while keeping the system message, and never cut the end of a response.
Two checks catch nearly every bug in this section, and they take a minute. Before a run, decode the trained tokens of a few dozen examples and read them. And log the fraction of examples with no trained tokens, which should be zero.
A loss mask from character offsets
Implement render_with_spans(messages) and labels_from_offsets(ids, offsets, spans).
render_with_spansrenders each message as"<|im_start|>" + role + "\n" + content + "<|im_end|>\n"and returns(text, spans), wherespanslists the(start, end)character range (end exclusive) of each assistant message's content plus its<|im_end|>, but not the newline after it.labels_from_offsetstakes token ids, their characteroffsets(a list of(start, end)pairs, end exclusive, with(0, 0)for tokens the tokenizer added that cover no characters, such as a BOS), and thespans. Return an integer NumPy array of unshifted labels (the Hugging Face convention), equal to the id where the token's character range overlaps some span by at least one character and elsewhere.
The real way
Here is the whole section in PyTorch. A small causal Transformer with learned position embeddings, three examples with masked labels, run once example by example and once packed with a block-diagonal mask and restarted positions. If packing is exact, the per-example losses must match. We also run the naive packing, with a plain causal mask, to see the contamination.
The packed losses equal the separate ones to every printed digit, and the naive packing changes every example after the first. With Hugging Face libraries, the template, the mask, and packing are configuration. apply_chat_template(..., return_assistant_tokens_mask=True, return_dict=True) returns a mask of assistant tokens for templates that mark them with {% generation %} blocks, and TRL's trainer wires the rest.
Read more
- Touvron et al. (2023). Llama 2: Open Foundation and Fine-Tuned Chat Models. Section 3.1 on SFT with prompt-token masking and packing, and Section 3.3 on Ghost Attention.
- Shi et al. (2024). Instruction Tuning With Loss Over Instructions. When keeping the loss on the prompt helps, and why.
- Huerta-Enochian and Ko (2024). Instruction Fine-Tuning: Does Prompt Loss Matter?. A study of fractional prompt loss weights.
- Lambert et al. (2024). Tulu 3: Pushing Frontiers in Open Language Model Post-Training. Section 4.3.2 on loss aggregation and the sum loss.
- Kundu et al. (2024). Enhancing Training Efficiency Using Packing with Flash Attention. Padding-free packing with position ids and variable-length attention in Hugging Face.
- Hugging Face. Chat templates. Templates, generation prompts, and assistant-token masks.
- OpenAI (2025). gpt-oss-120b and gpt-oss-20b Model Card. The harmony format, channels, and the instruction hierarchy.