Chat Templates and Loss Masking in Fine-Tuning

// In plain English

When you fine-tune a chat model, you hand it example conversations and say write replies like these. But a model doesn't see a tidy conversation with neat "User" and "Assistant" bubbles. It sees one long string of tokens. So two formatting details decide whether your training data actually teaches what you think it does: the chat template and the loss mask.

A chat template is the exact wrapping a model expects around each turn — the special tokens and role markers that say "a user message starts here" and "the assistant's reply starts here." Every model family has its own template, baked in during its own training. Use the wrong one and the model is reading a format it was never taught to recognize.

Loss masking decides which tokens the model is graded on while it learns. In a chat example, you only want it to learn how to produce the assistant's reply — not how to produce the user's question. Masking the loss means you tell the trainer: ignore the prompt tokens, only score the answer tokens. (You'll also hear this called train on completions only.)

// Why it matters

These two details share a nasty property: get them wrong and nothing crashes. Training runs to completion, the loss curve goes down, the run looks healthy — and then the model behaves slightly (or badly) wrong, and you have no obvious clue why. They are two of the most common silent fine-tuning bugs.

Why the chat template matters

A model learned during its own alignment to recognize one specific layout of role markers. That layout is part of how it "knows" where its turn begins and what voice to speak in. If you train on a different layout — say, you invent your own ### User: / ### Assistant: scheme, or you borrow another model's tokens — you fight against everything the base model already learned. Worse is a train/inference mismatch: you fine-tune with one template, then serve the model with a different one. At inference the model sees a frame it was never trained on for this task and produces garbled, run-on, or never-stopping output.

Why loss masking matters

If you train on every token — prompt included — you are spending part of the model's learning budget teaching it to predict the user's words. That's almost never what you want. You don't want the model to learn to write user questions; you want it to write great answers given those questions. Training on the prompt dilutes the signal, can make the model echo or hallucinate user-style text, and on tiny datasets it can meaningfully hurt quality. Masking focuses every gradient where it belongs: on the reply.

Both bugs hit hardest on small, high-quality datasets — exactly the regime most people fine-tune in. With a few hundred examples you can't afford to waste signal or confuse the model about its own format.

// How it works

Start from how a chat example becomes training data. You write conversations as a list of role/content messages. A chat template turns that list into one flat string of tokens by inserting the model's special markers. Then the trainer builds a label mask the same length as the token sequence, marking each token as either learn from this or ignore this.

What a chat template actually inserts

Different model families use different markers, but the idea is the same: tokens that name each role and tokens that signal where a turn ends. Here is the shape of a formatted two-turn conversation (illustrative, not any one model's exact tokens):

<|system|>
You are a helpful support agent.
<|end|>
<|user|>
What's the refund window?
<|end|>
<|assistant|>
Physical items can be returned within 30 days.
<|end|>

Those <|user|>, <|assistant|>, and end-of-turn markers are special tokens — single tokens in the model's vocabulary, not literal text it spells out letter by letter. That's why you can't just type your own role labels and expect the same behavior: the real markers are atoms the model was trained to recognize. The fix is to never hand-build this string. Use the tokenizer's built-in chat template (in the Hugging Face ecosystem, tokenizer.apply_chat_template(...)), which produces the exact format the model expects.

How the loss mask is built

After tokenizing, the trainer creates a labels array aligned with the tokens. Prompt tokens (system + user, plus the role markers leading into the assistant turn) get a special ignore value — in PyTorch the cross-entropy loss skips any label equal to -100. The assistant's reply tokens keep their real ids. The model still reads the whole prompt as context; it's just never graded on predicting those tokens.

Here's the masking idea in a few lines. Once the assistant turn's character span is known, anything before it becomes -100:

# token_ids: the full templated, tokenized conversation
# reply_start: index of the first ASSISTANT-reply token
#              (everything before it is prompt/markers)
IGNORE = -100
labels = token_ids.copy()
labels[:reply_start] = [IGNORE] * reply_start

# Cross-entropy with ignore_index=-100 skips the masked positions,
# so gradients flow ONLY from the assistant's reply tokens.
loss = cross_entropy(logits, labels, ignore_index=IGNORE)

// The classic template-mismatch bug

The single most common silent failure is a train/inference template mismatch. You fine-tune with one wrapping, then your serving stack applies a different one (or none). The model now sees a frame it doesn't recognize for this task, and the symptoms are weirdly specific:

• It never stops. The model keeps generating — answering its own follow-up questions, inventing a new user turn — because the end-of-turn token it learned isn't where it expects it.
• Wrong voice or role confusion. Replies drift into user-style phrasing, or the model restates the question instead of answering it.
• Format collapse on the first token. Output starts with stray markers or whitespace because the prompt didn't end exactly where training said the reply begins.
• Great eval loss, bad real output. Loss looked fine because it was measured with the training template; production uses another one.

A quick way to catch it

Before launching a run, print one fully-formatted, tokenized example and read it. Decode it back to text and confirm: the special tokens are real single tokens (not spelled-out characters), there's exactly one end-of-turn marker per turn, and the assistant reply begins right where your mask says it does. Five minutes of eyeballing here saves a wasted training run.

// Masking single-turn vs multi-turn chats

With one user turn and one assistant reply, masking is simple: ignore everything up to the reply, grade the reply. Multi-turn conversations need a small decision.

Grading all assistant turns is usually the right default — each reply is a valid training target, and you get more learning signal from one conversation. Grading only the last turn makes sense when earlier turns are fixed context you don't want the model to imitate. Either way, user turns are always masked out. The one mistake to avoid is masking the special tokens that end the assistant turn: the model must learn to emit its stop token, or it won't know when to stop at inference time.

// Going deeper

Once the basics click, a few finer points separate a clean fine-tune from a flaky one.

Tokenization boundaries are subtle. Whether a leading space or newline attaches to the prompt token or the first reply token can shift your mask by one position. An off-by-one mask either leaks a prompt token into the graded region or drops the reply's first token from it. This is exactly why you should trust the tokenizer's own template and offset mapping rather than splitting strings yourself.

Generation prompts. Templates often have an add_generation_prompt flag that appends the opening assistant marker (so the model knows it's its turn) without any reply text. You want that on at inference time and handled correctly at training time — the reply must start immediately after that marker, and your mask boundary must line up with it.

System prompts and packing. If you fine-tune with a particular system prompt, keep using it (or a compatible one) at inference. And if you pack multiple short conversations into one sequence for efficiency, make sure attention and masks don't let one example bleed into the next — every example needs its own clean prompt/reply boundary.

Does masking always help? For most supervised fine-tuning, completion-only masking is the safe default and often improves results, especially with little data. On very large, diverse datasets the gap can shrink. Treat it as a default you justify changing, not a knob to flip blindly — and confirm the effect when you evaluate the fine-tuned model.

It's mostly a data-prep concern. Templates and masks are decided when you build the dataset and configure the trainer, not in the model architecture. They sit right next to your other dataset preparation choices and apply whether you do a full fine-tune or a parameter-efficient one like LoRA — the mask logic is identical; only which weights update changes. Get these two details right and you remove a whole class of silent failures before they ever start.