5.2 Chat Templates, Data Formatting & Sequence Packing¶

Raw text pretraining teaches a model to continue arbitrary strings. Turning that capability into a conversational assistant requires convincing the model that conversations have structure: a system instruction, alternating user and assistant turns, and a definitive end-of-turn marker. This chapter is the engineering manual for that structure — how templates are designed, how loss masks are constructed to train only on completions, how multi-turn dialogues are packed tightly into fixed-length batches, and how tool calls extend the vocabulary of a chat turn.

Related chapters: Supervised Fine-Tuning & Instruction Tuning covers the overall SFT recipe; this chapter zooms into the data side. The Pretraining Objective & Loss explains the causal language modelling loss that we are selectively masking here. Tokenization: BPE, WordPiece, Unigram & Byte-Level explains how special tokens are added to a vocabulary. The capstone puts all of this to work: Post-Training: SFT, DPO, and Narrow RLVR (GRPO) That Works at 100M formats, masks and packs the Stack-100M SFT mix, and A Byte-Level BPE Tokenizer From Scratch reserves the chat control tokens at tokenizer-training time so the embedding table never has to be resized.

Why Chat Templates Exist¶

A pretrained language model has seen billions of documents, but almost none of them were labelled “this is the human’s turn; this is the assistant’s turn.” Without some delimiter, there is no way for the model—or the training code—to know where one turn ends and the next begins, and there is no way to enforce “generate here but not there” during both fine-tuning and inference.

The solution is a chat template: a deterministic function that maps a list of role–content pairs to a single token sequence that the model will both train on and generate from. Critically, a template also defines a prompt boundary: everything before this token is prompt; everything after is what the model is expected to generate. During training we mask (zero out the loss on) the prompt side and supervise only the generation side.

Templates must be:

1. Unambiguous — every special token must be one that the base model tokenizer will never split or confuse with natural text.
2. Consistent — inference must use the exact same string as training. A single missing space or wrong special token is enough to tank benchmark performance.
3. Efficient — the template overhead (role tokens, BOS/EOS) should be small relative to the content.

The Major Template Families¶

ChatML¶

ChatML (popularised by OpenAI and widely adopted thereafter) encodes each turn as:

<|im_start|>role\ncontent<|im_end|>\n

A complete two-turn exchange looks like this:

<|im_start|>system
You are a helpful assistant.<|im_end|>
<|im_start|>user
What is the capital of France?<|im_end|>
<|im_start|>assistant
Paris.<|im_end|>

The special tokens <|im_start|> and <|im_end|> are added to the base vocabulary during fine-tuning. ChatML — short for Chat Markup Language — was introduced by OpenAI alongside the gpt-3.5-turbo API in March 2023; the im in these tokens stands for instant message, i.e. a single chat turn, and has nothing to do with images. The role string (system, user, assistant) is ordinary text immediately followed by a newline; it is not itself a special token.

During generation, the server appends <|im_start|>assistant\n after the last user turn and lets the model decode until it emits <|im_end|>.

Llama 2 & Llama 3 Templates¶

Meta’s Llama 2 used a different convention:

[INST] <<SYS>>
{system}
<</SYS>>

{user_turn} [/INST] {assistant_turn} </s><s>[INST] {next_user} [/INST]

This is notoriously tricky: the system prompt is embedded inside the first [INST] block, there is no per-turn <s> except between turns, and the closing </s> is part of the assistant turn, not a standalone separator. Getting one of these details wrong silently produces bad outputs.

Llama 3 simplified things significantly, adopting a header-token system closer in spirit to ChatML:

<|begin_of_text|><|start_header_id|>system<|end_header_id|>

You are a helpful assistant.<|eot_id|><|start_header_id|>user<|end_header_id|>

What is the capital of France?<|eot_id|><|start_header_id|>assistant<|end_header_id|>

Paris.<|eot_id|>

Here <|start_header_id|>, <|end_header_id|>, and <|eot_id|> (end-of-turn) are proper special tokens. The double newline after the header acts as a role separator. <|eot_id|> replaces both [/INST] and </s> in a single unambiguous token.

Qwen Templates¶

Qwen models (Alibaba’s open-weight series) use a format almost identical to ChatML, with <|im_start|> / <|im_end|> inherited from OpenAI conventions:

<|im_start|>system
You are Qwen, created by Alibaba Cloud.<|im_end|>
<|im_start|>user
Translate "bonjour" to English.<|im_end|>
<|im_start|>assistant
Hello.<|im_end|>

The Qwen tokenizer’s special-token list includes <|im_start|>, <|im_end|>, <|endoftext|>, and a set of tool-calling tokens discussed later.

Comparison¶

The HuggingFace tokenizers library stores a Jinja2 template string in the tokenizer’s chat_template field so that tokenizer.apply_chat_template(messages) always returns the correct string for that model. This is the canonical way to apply templates in Python — do not hard-code delimiter strings yourself.

Building a Chat Template from Scratch¶

Let us implement a minimal ChatML formatter and a corresponding loss-mask builder. This is the kind of code you would write to process a dataset before passing it to a training loop.

"""
chat_template.py
────────────────
A from-scratch ChatML formatter with a loss mask builder.
Compatible with any tokenizer that has <|im_start|> and <|im_end|>
in its vocabulary (e.g. Qwen, phi-3, many fine-tuned models).
"""

from __future__ import annotations
from dataclasses import dataclass
from typing import Literal, Sequence
import torch
from transformers import PreTrainedTokenizerFast

# ---------------------------------------------------------------------------
# Data structures
# ---------------------------------------------------------------------------

Role = Literal["system", "user", "assistant", "tool"]

@dataclass
class ChatMessage:
    role: Role
    content: str

# ---------------------------------------------------------------------------
# Template formatter
# ---------------------------------------------------------------------------

IM_START = "<|im_start|>"
IM_END   = "<|im_end|>"

def format_chatml(
    messages: Sequence[ChatMessage],
    add_generation_prompt: bool = True,
) -> str:
    """
    Render a list of ChatMessage objects to a ChatML string.

    If add_generation_prompt=True (used during inference and at the
    end of training examples), appends '<|im_start|>assistant\n' so
    the model knows it should generate next.
    """
    pieces: list[str] = []
    for msg in messages:
        # Each turn: <|im_start|>ROLE\nCONTENT<|im_end|>\n
        pieces.append(f"{IM_START}{msg.role}\n{msg.content}{IM_END}\n")

    if add_generation_prompt:
        pieces.append(f"{IM_START}assistant\n")

    return "".join(pieces)


# ---------------------------------------------------------------------------
# Loss mask builder
# ---------------------------------------------------------------------------

def build_chatml_loss_mask(
    tokenizer: PreTrainedTokenizerFast,
    messages: Sequence[ChatMessage],
    max_length: int = 2048,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Tokenize a chat conversation and return:
        input_ids  : LongTensor of shape (seq_len,)
        loss_mask  : BoolTensor  of shape (seq_len,)
                     True  → compute loss on this position
                     False → mask out (prompt / special tokens)

    Strategy: we tokenize the full sequence, then for each assistant
    turn we find the span [start_of_content, end_of_turn] and set
    loss_mask = True only for those positions.
    """
    # Render full string (no generation prompt needed for training)
    full_text = format_chatml(messages, add_generation_prompt=False)

    # Tokenize full sequence; do not add special tokens automatically
    # (the template already handles BOS/EOS via <|im_start|> etc.)
    enc = tokenizer(
        full_text,
        return_tensors="pt",
        truncation=True,
        max_length=max_length,
        add_special_tokens=False,
    )
    input_ids: torch.Tensor = enc["input_ids"][0]  # (seq_len,)
    seq_len = input_ids.shape[0]

    loss_mask = torch.zeros(seq_len, dtype=torch.bool)

    # -----------------------------------------------------------------------
    # Find assistant spans by scanning for the token IDs of
    # <|im_start|>assistant\n  …content…  <|im_end|>
    # -----------------------------------------------------------------------
    im_start_id: int = tokenizer.convert_tokens_to_ids(IM_START)
    im_end_id:   int = tokenizer.convert_tokens_to_ids(IM_END)

    # Tokenize the role string that follows <|im_start|> for an assistant turn.
    # We need to detect "assistant\n" as token(s).
    role_ids: list[int] = tokenizer.encode(
        "assistant\n", add_special_tokens=False
    )

    ids = input_ids.tolist()
    i = 0
    while i < len(ids):
        # Look for <|im_start|>
        if ids[i] != im_start_id:
            i += 1
            continue

        # Check if the next tokens are the assistant role
        role_end = i + 1 + len(role_ids)
        if role_end <= len(ids) and ids[i + 1 : role_end] == role_ids:
            # We are inside an assistant turn.
            # Content starts at role_end; supervise until <|im_end|>.
            content_start = role_end
            # Find the matching <|im_end|>
            j = content_start
            while j < len(ids) and ids[j] != im_end_id:
                j += 1
            # Mark content tokens for loss (include <|im_end|> itself so the
            # model learns to emit the stop token)
            if j < len(ids):
                loss_mask[content_start : j + 1] = True
            i = j + 1
        else:
            # Not an assistant turn; skip to the next <|im_end|>
            i += 1

    return input_ids, loss_mask

messages = [
    ChatMessage("system",    "You are a helpful assistant."),
    ChatMessage("user",      "What is 2+2?"),
    ChatMessage("assistant", "4."),
]

<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n
<|im_start|>user\nWhat is 2+2?<|im_end|>\n
<|im_start|>assistant\n4.<|im_end|>\n

Shipping It: The Jinja Template and the New Special Tokens¶

A hand-written scanner like the one above is the right way to understand masking, but the artefact you ship is a Jinja2 string stored on the tokenizer, because that is what apply_chat_template — and every serving stack that reads tokenizer_config.json (vLLM, SGLang, TGI, llama.cpp) — will execute. Transformers extends Jinja with a {% generation %} block that marks precisely the spans the assistant is expected to produce, so the library can hand you the loss mask instead of you scanning for it:

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tok   = AutoTokenizer.from_pretrained("path/to/stack-100m-base")
model = AutoModelForCausalLM.from_pretrained("path/to/stack-100m-base")

# 1. Add the ChatML control tokens. (Skip if the tokenizer was trained with
#    them already reserved — see the capstone tokenizer chapter.)
n_new = tok.add_special_tokens(
    {"additional_special_tokens": ["<|im_start|>", "<|im_end|>"]}
)
tok.eos_token = "<|im_end|>"      # generation must stop at end-of-turn

# 2. Attach the template. Each piece is a RAW string so that the "\n" stays a
#    two-character Jinja escape instead of becoming a literal newline inside a
#    Jinja string literal (which is a parse error).
tok.chat_template = (
    r"{% for m in messages %}"
    r"{% if m['role'] == 'assistant' %}"
    r"{{ '<|im_start|>assistant\n' }}"
    r"{% generation %}{{ m['content'] + '<|im_end|>' }}{% endgeneration %}"
    r"{{ '\n' }}"
    r"{% else %}"
    r"{{ '<|im_start|>' + m['role'] + '\n' + m['content'] + '<|im_end|>\n' }}"
    r"{% endif %}"
    r"{% endfor %}"
    r"{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"
)
tok.save_pretrained("path/to/stack-100m-sft")   # template travels with weights

# 3. Now the library builds the mask for you — no scanner required.
out = tok.apply_chat_template(
    [{"role": "user", "content": "What is 2+2?"},
     {"role": "assistant", "content": "4."}],
    tokenize=True,
    return_dict=True,
    return_assistant_tokens_mask=True,   # requires the {% generation %} block
)
input_ids = torch.tensor(out["input_ids"])
loss_mask = torch.tensor(out["assistant_masks"], dtype=torch.bool)

New special tokens also need embeddings. Rows appended by resize_token_embeddings are randomly initialised and, because each new token appears at most a handful of times per batch, they can stay near that random init for a long time — which is why <|im_end|> is sometimes emitted unreliably early in a run. Initialising them at the mean of the existing embedding rows starts them in-distribution and is a one-line fix:

old_vocab = model.get_input_embeddings().weight.shape[0]
model.resize_token_embeddings(len(tok))          # new rows are appended at the end
with torch.no_grad():
    for mod in (model.get_input_embeddings(), model.get_output_embeddings()):
        if mod is None:                          # some models have no separate head
            continue
        W = mod.weight
        W[old_vocab:] = W[:old_vocab].mean(dim=0, keepdim=True)

For Stack-100M you can sidestep this entirely: reserve <|im_start|>, <|im_end|> and a block of spare <|reserved_N|> ids when you train the BPE tokenizer, so the vocabulary — and therefore every checkpoint’s embedding shape — is fixed from the first pretraining step. The rows are still effectively untrained (those ids never occur in raw web text), so the mean-init trick above is still worth applying at the start of SFT.

Loss Masking in Depth¶

Why Mask the Prompt?¶

During causal language modelling pretraining, every token in the sequence is both an input and a target. When we fine-tune on conversations, we do not want the model to “memorise” system prompts or user questions as outputs it should produce; we want it to learn to respond to them. Including the prompt in the loss does two things:

1. Gradient pollution: the gradient signal for role tokens like <|im_start|>user is dominated by examples where those tokens appear, and none of that signal helps the model become a better assistant.
2. Objective mismatch: at inference time, the model never generates the user’s message; training it to do so creates a distribution mismatch.

Empirically, SFT runs where the prompt is masked converge to lower perplexity on held-out completions (not full sequences) and exhibit fewer repetitions of the system prompt in responses.

The Loss Computation with a Mask¶

The standard causal LM loss for a sequence \(x_1, \ldots, x_T\) is:

With a binary loss mask \(m_t \in \{0, 1\}\), we change this to:

The denominator normalises by the number of supervised tokens rather than the total sequence length, so that longer prompts do not dilute the gradient.

That denominator hides a subtlety once you add gradient accumulation or data parallelism. If every micro-batch divides by its own supervised-token count and you then average the micro-batch losses, each micro-batch gets equal weight no matter how much signal it carried: a micro-batch with 20 assistant tokens moves the update as much as one with 2,000, and the result is no longer the token-mean over the global batch. The fix is to accumulate unnormalised token losses and divide once by the global supervised-token count. HuggingFace Trainer does this by threading a num_items_in_batch argument down into the model’s loss function, and average_tokens_across_devices (in TrainingArguments, inherited by TRL’s SFTConfig) all-reduces that count across data-parallel ranks. Packing makes this matter more, not less, because packed rows vary widely in how many supervised tokens they contain.

In PyTorch the standard way is to pass labels to a causal LM where masked positions are set to -100 — CrossEntropyLoss ignores index -100 by default.

def labels_from_mask(
    input_ids: torch.Tensor,
    loss_mask: torch.Tensor,
) -> torch.Tensor:
    """
    Construct labels for a causal LM: supervised positions keep their
    token id; masked positions are set to -100 (ignored by CrossEntropyLoss).

    Note the standard causal LM shift: the model predicts x_{t+1} given x_{<=t},
    so labels are input_ids shifted left by one, i.e. labels[t] = input_ids[t+1].
    HuggingFace CausalLM models do this shift internally when you pass `labels`,
    so we just pass the full input_ids with -100 at masked positions.
    """
    labels = input_ids.clone()
    labels[~loss_mask] = -100
    return labels

Multi-Turn Loss Masking¶

In a multi-turn conversation there are multiple assistant turns, each of which should be supervised independently. The loss-mask builder above already handles this: it scans the entire token sequence and marks all assistant-content spans, not just the final one.

A common mistake is to supervise only the last assistant turn. This wastes signal — a four-turn conversation has four assistant responses, all of which are valuable training signal.

The causal mask of the transformer still lets each assistant token attend to everything before it (including previous user turns), so learning is coherent — the model sees the full context, it just does not receive gradient for repeating the prompt tokens.

System Prompts and Role Tokens¶

The System Turn¶

The system prompt is the privileged message that gives the model its persona, its capabilities, and its constraints. It appears as a system role turn at the very beginning of the context. During training on conversations that include system prompts, the system turn is part of the prompt and is masked from the loss.

Not all models have a dedicated system role. Llama 2’s template embeds the system prompt inside the first [INST] block. Mistral v1 has no system slot at all and expects operators to prepend the system content to the first user message. This matters for fine-tuning: if you are training on Mistral but your dataset has a role: system field, you must decide whether to drop it, inject it into the user turn, or add <s>[INST] tokens — and then be consistent at inference time.

Role Tokens as Special Tokens¶

When role markers like <|im_start|> are added to the tokenizer vocabulary, they receive their own embedding vectors in the model’s embedding table. These are randomly initialised at fine-tuning time (unless the base model already included them during pretraining, as many modern bases do). Because only a handful of examples per token appear in each gradient step, the embeddings for <|im_end|> and role names typically need a higher learning rate or more warm-up steps than ordinary parameters, or they remain near their random init.

One practical consequence: if you fine-tune a model on ChatML format but then serve it with a Llama 3-format prompt, the special tokens your model was trained with may not exist in the serving tokenizer — and vice versa. Always version and lock your tokenizer alongside your model weights.

Sequence Packing¶

The Problem with Padding¶

Naive batching left-pads (or right-pads) every sequence in a batch to the same length, filling the gap with a [PAD] token. The model then wastes computation on padding positions. For SFT datasets where individual examples range from 50 to 2,000 tokens, a batch max-length of 2,048 means that short examples may waste 90%+ of their sequence slot on padding.

For a 7B-parameter model with context length 2,048, each forward + backward pass over a batch of size 4 costs on the order of:

If the average sequence is only 512 tokens long, padding wastes 75% of those FLOPs. Sequence packing eliminates this waste.

Pack-Then-Truncate¶

Sequence packing concatenates multiple independent examples into a single sequence of exactly max_length tokens. The simplest algorithm is first-fit-decreasing bin packing: sort examples by length, greedily fill each bin.

"""
packing.py
──────────
First-fit-decreasing bin packer for SFT sequences.
Each bin becomes one training row of length `max_length`.
"""

from __future__ import annotations
from typing import NamedTuple
import torch

class PackedSample(NamedTuple):
    input_ids:  torch.Tensor  # (max_length,)
    labels:     torch.Tensor  # (max_length,) with -100 for masked positions
    # position_ids tracks the per-document position so RoPE is applied
    # correctly within each packed document.
    position_ids: torch.Tensor  # (max_length,)
    # seqlens holds the length of each real (non-pad) document in this bin,
    # so the collator can build cu_seqlens for Flash Attention's varlen path.
    seqlens: torch.Tensor  # (num_docs,) int32

PAD_ID = 0   # tokenizer.pad_token_id — used only to fill the last bin

def pack_sequences(
    examples: list[tuple[torch.Tensor, torch.Tensor]],  # (input_ids, labels)
    max_length: int = 2048,
    pad_id: int = PAD_ID,
) -> list[PackedSample]:
    """
    Pack variable-length (input_ids, labels) pairs into fixed-length bins.

    For each packed bin we also build position_ids that reset to 0 at the
    start of each new document — critical for correct RoPE / ALiBi behaviour.
    """
    # Sort longest-first for better bin utilisation (first-fit-decreasing)
    examples = sorted(examples, key=lambda x: x[0].shape[0], reverse=True)

    bins: list[list[tuple[torch.Tensor, torch.Tensor]]] = []
    bin_lengths: list[int] = []

    for ids, labs in examples:
        L = ids.shape[0]
        if L > max_length:
            # Truncate oversized examples
            ids  = ids[:max_length]
            labs = labs[:max_length]
            L = max_length

        placed = False
        for b_idx, b_len in enumerate(bin_lengths):
            if b_len + L <= max_length:
                bins[b_idx].append((ids, labs))
                bin_lengths[b_idx] += L
                placed = True
                break
        if not placed:
            bins.append([(ids, labs)])
            bin_lengths.append(L)

    packed: list[PackedSample] = []
    for bin_contents, bin_len in zip(bins, bin_lengths):
        # Concatenate all examples in this bin
        all_ids  = torch.cat([x[0] for x in bin_contents])
        all_labs = torch.cat([x[1] for x in bin_contents])

        # Build position_ids: reset to 0 at the start of each document
        pos = []
        for ids, _ in bin_contents:
            pos.append(torch.arange(ids.shape[0]))
        position_ids = torch.cat(pos)

        # Pad the bin to max_length if needed
        pad_len = max_length - bin_len
        if pad_len > 0:
            all_ids  = torch.cat([all_ids,  torch.full((pad_len,), pad_id)])
            all_labs = torch.cat([all_labs, torch.full((pad_len,), -100)])
            position_ids = torch.cat([position_ids, torch.zeros(pad_len, dtype=torch.long)])

        # Per-document lengths within this bin (real docs only; the trailing
        # pad region is handled by the collator via max_length - content_len).
        doc_lens = torch.tensor(
            [ids.shape[0] for ids, _ in bin_contents], dtype=torch.int32
        )
        packed.append(PackedSample(all_ids, all_labs, position_ids, doc_lens))

    return packed

Attention Boundaries: Preventing Cross-Document Leakage¶

The most important correctness issue with packing is cross-document attention leakage. In a naively packed sequence, token 1,024 (the first token of document 2) can attend to token 1,023 (the last token of document 1), creating a spurious dependency that never exists at inference time.

For most SFT training this leakage is a minor nuisance — the model quickly learns that <|im_start|> marks a boundary — but it can degrade performance on short conversations and is measurably harmful for preference learning where sequence independence is critical.

The fix is a block-diagonal attention mask: within each packed bin, the causal mask is restricted to within-document boundaries. Each document can only attend to its own prior tokens.

def make_block_diagonal_mask(
    document_lengths: list[int],
    dtype: torch.dtype = torch.float32,
    neg_inf: float = float("-inf"),
) -> torch.Tensor:
    """
    Build a causal block-diagonal attention mask for a packed sequence.

    Returns a (total_len, total_len) additive mask M where
        M[i, j] = 0        if j <= i AND same document as i
        M[i, j] = -inf     otherwise

    Add this to raw attention logits before softmax.
    """
    total_len = sum(document_lengths)
    # Start with full -inf (no attention allowed anywhere)
    mask = torch.full((total_len, total_len), neg_inf, dtype=dtype)

    offset = 0
    for doc_len in document_lengths:
        # Within this document, causal attention is allowed
        for i in range(doc_len):
            for j in range(i + 1):  # j <= i → causal
                mask[offset + i, offset + j] = 0.0
        offset += doc_len

    return mask  # (total_len, total_len)

In practice, building this mask explicitly is memory-intensive for long sequences: it is \(O(T^2)\) bytes, so a single 8,192-token packed row costs 268 MB in fp32 — larger than the activations it guards. Production implementations therefore never materialise it. Transformers’ DataCollatorWithFlattening (the collator behind TRL’s padding_free=True) concatenates a micro-batch into a single row of shape (1, total_tokens) and emits position_ids that reset to 0 at each document boundary; the Flash-Attention backend recovers cu_seqlens from those resets and dispatches to the varlen kernel, so document isolation comes for free. Megatron-LM does the equivalent by handing an explicit PackedSeqParams (cumulative sequence lengths plus max segment length) to its TransformerEngine fused-attention path. In both cases the boundary information travels as an \(O(\text{num docs})\) offset vector rather than a \(T \times T\) matrix.

# Using Flash Attention's varlen API (pseudo-code showing the key argument)
# flash_attn.flash_attn_varlen_func accepts:
#   qkv : packed Q/K/V for all documents concatenated
#   cu_seqlens_q, cu_seqlens_k : cumulative sequence lengths (int32)
#   max_seqlen_q, max_seqlen_k : max document length in the batch
# This naturally restricts attention to within-document boundaries.
import flash_attn

output = flash_attn.flash_attn_varlen_func(
    q=q_packed,            # (total_tokens, n_heads, head_dim)
    k=k_packed,
    v=v_packed,
    cu_seqlens_q=cu_seqlens,  # e.g. tensor([0, 512, 1024, 2048])
    cu_seqlens_k=cu_seqlens,
    max_seqlen_q=max_doc_len,
    max_seqlen_k=max_doc_len,
    causal=True,
)

Tool-Call Formatting¶

Modern assistant models must handle not just text exchanges but also tool calls — structured requests to invoke external functions — and their results. The chat template is the natural place to encode these, and several distinct conventions have emerged.

The Tool-Use Turn¶

A tool-call turn is an assistant message that contains a structured payload instead of (or in addition to) plain text. The template must encode both the call and the return value. In ChatML-based systems, this is often done with a dedicated tool role for the tool’s response:

[
  {"role": "user",      "content": "What's the weather in Paris?"},
  {"role": "assistant", "content": null,
   "tool_calls": [{"id": "call_1", "type": "function",
                   "function": {"name": "get_weather",
                                "arguments": "{\"city\": \"Paris\"}"}}]},
  {"role": "tool",      "content": "{\"temp\": 18, \"condition\": \"cloudy\"}",
   "tool_call_id": "call_1"},
  {"role": "assistant", "content": "It's 18 °C and cloudy in Paris."}
]

When this is rendered to tokens, the template serialises the tool_calls JSON (or a subset of it) as the assistant turn’s content. Some models (Qwen, Llama 3.1+) define dedicated special tokens for the open/close of a tool call block:

<|im_start|>assistant
<tool_call>
{"name": "get_weather", "arguments": {"city": "Paris"}}
</tool_call><|im_end|>
<|im_start|>tool
{"temp": 18, "condition": "cloudy"}<|im_end|>
<|im_start|>assistant
It's 18 °C and cloudy in Paris.<|im_end|>

Here <tool_call> and </tool_call> may be single special tokens or ordinary text (model-dependent). The key design decision is whether they are special tokens — in which case they are immune to BPE splitting — or ordinary strings that must be stable under tokenization.

Loss Masking for Tool Calls¶

For training, we must decide which parts of a tool-call exchange to supervise:

• The assistant’s tool-call invocation (the JSON payload): supervised — the model must learn to emit valid JSON with the right function name and arguments.
• The tool role turn (the function return): masked — the model does not generate this; the harness does.
• The final assistant text after the tool return: supervised.

The loss-mask scanner from our build_chatml_loss_mask function handles this correctly if tool turns are treated as non-assistant turns (they will not match the role_ids == "assistant\n" check).

"""
Tool-call example: extend build_chatml_loss_mask to recognise
the <tool_call>...</tool_call> block within an assistant turn.

For simplicity we supervise the ENTIRE assistant turn content
(including the tool_call JSON), since the model should learn to
emit both plain text and tool call syntax.
"""

# The loss masking logic in build_chatml_loss_mask already handles this:
# it supervises everything between <|im_start|>assistant\n and <|im_end|>,
# which includes any <tool_call> blocks the model emits.
# The only special case is if you want to EXCLUDE the <tool_call> JSON
# from the loss (e.g. treat it as planning overhead).  For most SFT
# practitioners, supervising it is the right choice.

Putting It All Together: A Training-Ready Collator¶

Here is a complete data collator that formats, packs, and masks a batch of conversations, ready to be passed to a training loop.

"""
collator.py
───────────
A DataCollator that applies ChatML formatting, loss masking, and
optional sequence packing.  Drop this into any HuggingFace Trainer.
"""

from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any
import torch
from transformers import PreTrainedTokenizerFast

# Re-use helpers defined above
# from chat_template import build_chatml_loss_mask, labels_from_mask, ChatMessage
# from packing import pack_sequences

@dataclass
class ChatMLCollator:
    tokenizer: PreTrainedTokenizerFast
    max_length: int = 2048
    pack: bool = True

    def __call__(self, raw_batch: list[dict[str, Any]]) -> dict[str, torch.Tensor]:
        """
        raw_batch: list of dicts with key "messages": list[dict] with
                   "role" and "content" keys (HuggingFace-style).
        """
        examples: list[tuple[torch.Tensor, torch.Tensor]] = []

        for item in raw_batch:
            # Convert raw dicts to ChatMessage objects
            messages = [
                ChatMessage(role=m["role"], content=m["content"])
                for m in item["messages"]
            ]
            # Build tokens + mask
            input_ids, loss_mask = build_chatml_loss_mask(
                self.tokenizer, messages, max_length=self.max_length
            )
            labels = labels_from_mask(input_ids, loss_mask)
            examples.append((input_ids, labels))

        if self.pack:
            packed = pack_sequences(
                examples,
                max_length=self.max_length,
                pad_id=self.tokenizer.pad_token_id or 0,
            )
            batch_ids  = torch.stack([p.input_ids    for p in packed])
            batch_labs = torch.stack([p.labels       for p in packed])
            batch_pos  = torch.stack([p.position_ids for p in packed])

            # Build a length-based attention_mask and cu_seqlens over the
            # flattened (B * max_length,) batch. The trailing pad region of
            # each row becomes its own segment so offsets sum to max_length.
            attention_mask = torch.zeros_like(batch_ids)
            row_seqlens: list[torch.Tensor] = []
            for r, p in enumerate(packed):
                content_len = int(p.seqlens.sum())          # real (non-pad) tokens
                attention_mask[r, :content_len] = 1
                lens = p.seqlens
                pad_len = self.max_length - content_len
                if pad_len > 0:
                    lens = torch.cat(
                        [lens, torch.tensor([pad_len], dtype=torch.int32)]
                    )
                row_seqlens.append(lens)

            flat_seqlens = torch.cat(row_seqlens)           # (num_segments_total,)
            cu_seqlens = torch.zeros(flat_seqlens.numel() + 1, dtype=torch.int32)
            torch.cumsum(flat_seqlens, dim=0, out=cu_seqlens[1:])
            max_seqlen = int(flat_seqlens.max())
        else:
            # Simple right-padding without packing
            max_len = max(ids.shape[0] for ids, _ in examples)
            pad_id  = self.tokenizer.pad_token_id or 0

            def pad(t: torch.Tensor, fill: int) -> torch.Tensor:
                p = max_len - t.shape[0]
                return torch.cat([t, torch.full((p,), fill)])

            batch_ids  = torch.stack([pad(ids,  pad_id) for ids, _    in examples])
            batch_labs = torch.stack([pad(labs, -100)   for _,   labs in examples])
            # Standard sequential positions (no packing, so no reset needed)
            batch_pos  = torch.arange(max_len).unsqueeze(0).expand(len(examples), -1)

            # Length-based attention_mask. Do NOT use (batch_ids != pad_token_id):
            # when pad_token == eos_token (the Llama-style config from 5.1) that
            # would also zero out every real EOS / <|im_end|> token.
            attention_mask = torch.zeros_like(batch_ids)
            for r, (ids, _) in enumerate(examples):
                attention_mask[r, : ids.shape[0]] = 1
            cu_seqlens = None
            max_seqlen = None

        batch = {
            "input_ids":      batch_ids,
            "labels":         batch_labs,
            "attention_mask": attention_mask,
            "position_ids":   batch_pos,
        }
        # cu_seqlens / max_seqlen drive the flash-attn varlen path (see below)
        # that keeps attention inside document boundaries. A stock HF model does
        # not accept these kwargs, so consume them in a patched attention forward
        # (or a custom Trainer) and pop them before calling model().
        if self.pack:
            batch["cu_seqlens"] = cu_seqlens
            batch["max_seqlen"] = max_seqlen
        return batch

Wire the packed batch into Flash Attention’s variable-length kernel inside your (patched) attention forward:

# q, k, v : (B, L, n_heads, head_dim) for the packed batch. Flatten the (B, L)
# dims and hand Flash Attention the cu_seqlens the collator computed.
import flash_attn

B, L, H, D = q.shape
out = flash_attn.flash_attn_varlen_func(
    q.reshape(B * L, H, D),
    k.reshape(B * L, H, D),
    v.reshape(B * L, H, D),
    cu_seqlens_q=batch["cu_seqlens"],   # int32, ends at B * max_length
    cu_seqlens_k=batch["cu_seqlens"],
    max_seqlen_q=batch["max_seqlen"],
    max_seqlen_k=batch["max_seqlen"],
    causal=True,
)  # attention is now confined to each document; no cross-doc leakage

Verify it. With pack=True, the offsets and mask must be self-consistent: assert int(batch["cu_seqlens"][-1]) == batch["input_ids"].numel() (segments cover every token) and assert (batch["labels"][batch["attention_mask"] == 0] == -100).all() (every masked-out position is unsupervised). For the pad==eos trap, build a batch with a tokenizer where pad_token_id == eos_token_id and confirm attention_mask.sum() equals the number of real tokens — not real-tokens-minus-the-EOS-tokens, which is what the old (batch_ids != pad_token_id) mask would have produced.

The Library Path: TRL’s SFTTrainer¶

You now know exactly what has to happen, which means you are qualified to let a maintained library do it. TRL’s SFTTrainer performs chat-template rendering, completion-only masking, bin packing and padding-free attention from configuration flags alone:

from trl import SFTTrainer, SFTConfig

# Rows of `train_ds` look like {"messages": [{"role": ..., "content": ...}, ...]};
# SFTTrainer calls tokenizer.apply_chat_template on them for you.
trainer = SFTTrainer(
    model="path/to/stack-100m-base",
    train_dataset=train_ds,
    args=SFTConfig(
        max_length=2048,
        packing=True,                        # bin-pack examples into max_length rows
        assistant_only_loss=True,            # needs {% generation %} in the template
        average_tokens_across_devices=True,  # global supervised-token normaliser
        bf16=True,
    ),
)
trainer.train()

Each flag maps onto something we built by hand: assistant_only_loss=True is build_chatml_loss_mask (it consumes the assistant_masks that the {% generation %} block produces), and packing=True is pack_sequences plus the varlen call. Recent TRL versions pack with resetting position_ids so the Flash-Attention backend isolates documents; with packing=False you get the same varlen benefit on a plain right-padded batch via padding_free=True. Confirm the behaviour on your installed version rather than assuming it — the from-scratch code in this chapter is exactly what you fall back on to audit a suspicious run. The Stack-100M SFT stage in Post-Training: SFT, DPO, and Narrow RLVR (GRPO) That Works at 100M uses this configuration; its tool-calling data follows the format in A Narrow Auto-Research Agent.

Practical Checklist for Template Consistency¶

Template bugs are among the most common causes of unexplained degradation in SFT runs. The following checklist covers the most frequent footguns.

1. Lock tokenizer version. A tokenizer update can reorder or renumber special tokens. Pin your tokenizer to the same commit as your model weights.

2. Use apply_chat_template, not string concatenation. tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False) is always right; manual string formatting is often subtly wrong. If your template carries {% generation %} blocks, add return_dict=True, return_assistant_tokens_mask=True and take the mask from the library rather than re-deriving it.

3. Verify the loss mask before training. Print out a sample batch and confirm that labels != -100 only for assistant content. A one-liner: (labels != -100).float().mean() should match your expected supervised-fraction (typically 0.25–0.50 for chat data).

4. Match the inference prompt exactly. The string passed to the model server at inference time must be the output of apply_chat_template(messages, add_generation_prompt=True). Many bugs arise from adding an extra space, a missing BOS token, or the wrong role name.

5. Check EOS handling. Most templates include an EOS-like token (<|im_end|>, <|eot_id|>) at the end of each assistant turn. Make sure your inference server stops generation at this token — add it to stop_sequences if it is not the model’s default EOS.

6. Truncate from the left for long prompts. If a conversation exceeds max_length, truncating from the right destroys the final assistant turn. Truncate early user turns from the left, preserving the system prompt and the most recent turns.

7. Tool-turn masking. If your dataset includes tool-return turns (role: "tool"), confirm that these are masked. The model should not learn to generate tool return values.

Cross-reference: Supervised Fine-Tuning & Instruction Tuning for the broader training loop; PEFT I: LoRA, QLoRA, DoRA & The Adapter Family for parameter-efficient fine-tuning considerations; Tool Use & Function Calling for the inference-side view of tool calls.

Further Reading¶

• HuggingFace chat_templates documentation and Jinja2 template examples — the canonical reference for apply_chat_template and per-model template strings. Available in the transformers repository under docs/source/en/chat_templating.md.
• Meta, “Llama 2: Open Foundation and Fine-Tuned Chat Models” (Touvron et al., 2023) — describes the [INST]/[/INST] format and its role in RLHF data collection.
• Meta, “The Llama 3 Herd of Models” (Dubey et al., 2024) — introduces the Llama 3 header-token template and discusses multi-turn RLHF data formatting.
• OpenAI, “ChatML format” — the original description of <|im_start|>/<|im_end|> tokens, introduced with the gpt-3.5-turbo API in 2023 in OpenAI’s openai-python ChatML documentation (chatml.md).
• Krell et al., “Efficient Sequence Packing without Cross-contamination” (2021) — empirical study on the impact of cross-document attention in packed pretraining sequences, and the efficacy of document-boundary masks.
• Tri Dao, “FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning” (2023) — the flash_attn_varlen_func API used for packing without explicit mask materialisation.
• Zheng et al., “SGLang: Efficient Execution of Structured Language Model Programs” (2023) — covers structured generation and how tool-call schemas interact with constrained decoding, complementing the template story.

Exercises¶

1. (Conceptual) The chapter insists that during SFT we set labels = -100 on every prompt and role token and supervise only the assistant-turn content. Suppose instead you trained on the full sequence — every token both an input and a target, exactly as in pretraining. Name the two concrete failure modes the chapter attributes to this choice, and explain what each one does to the resulting model.

2. (Quantitative) Take the worked example from the chapter: a 26-token packed conversation in which exactly 3 tokens are assistant content (loss_mask = True) and the other 23 are masked. During one forward pass the model assigns the following per-token negative log-probabilities \(-\log p_\theta(x_t \mid x_{<t})\) to the three supervised positions: \(0.50\), \(1.00\), \(0.30\). Using the chapter’s masked-loss formula, compute \(\mathcal{L}_{\text{masked}}\). Then compute what the loss would have been if you had (incorrectly) normalised by the full sequence length \(T = 26\) while still summing only the supervised terms. By what factor do the two differ, and which normaliser does the chapter say is correct?

3. (Quantitative) You are packing an SFT dataset with max_length = 1024. The five examples have token lengths: A = 700, B = 500, C = 450, D = 300, E = 150. Apply the chapter’s first-fit-decreasing bin packer (sort longest-first, then place each example in the first bin it fits). List the resulting bins and their fill. Then compute the token utilisation both without packing (one right-padded row per example) and with packing, and state how many training rows each approach produces.

4. (Conceptual) In the packed collator, a colleague sets pack=True and builds the attention mask with the one-liner attention_mask = (batch_ids != tokenizer.pad_token_id). Their tokenizer follows the Chapter 5.1 Llama-style convention pad_token = eos_token. Describe the two independent things wrong here — one specific to the pad == eos choice, and one intrinsic to using any 2D row-level mask with packing — and state the chapter’s fix for each.

5. (Implementation) The chapter’s build_chatml_loss_mask supervises the entire assistant turn, including any <tool_call>...</tool_call> JSON. Some practitioners prefer to treat the tool-call JSON as “planning overhead” and exclude it from the loss while still supervising the surrounding assistant prose. Write a function mask_tool_call_json(input_ids, loss_mask, tokenizer) that takes the tensors returned by build_chatml_loss_mask and sets loss_mask = False on every token from <tool_call> through </tool_call> inclusive. Assume <tool_call> and </tool_call> are ordinary text that tokenizes to a stable id sequence (not single special tokens). Keep it consistent with the chapter’s scanning style.

import torch
from transformers import PreTrainedTokenizerFast

def _find_subseq(ids: list[int], pat: list[int], start: int = 0) -> int:
    """Return the index of the first occurrence of pat in ids at/after
    `start`, or -1 if not found."""
    if not pat:
        return -1
    last = len(ids) - len(pat)
    for i in range(start, last + 1):
        if ids[i : i + len(pat)] == pat:
            return i
    return -1

def mask_tool_call_json(
    input_ids: torch.Tensor,   # (seq_len,)
    loss_mask: torch.Tensor,   # (seq_len,) bool, from build_chatml_loss_mask
    tokenizer: PreTrainedTokenizerFast,
) -> torch.Tensor:
    """
    Set loss_mask = False across every <tool_call> ... </tool_call> span
    (inclusive of both delimiters), so the tool-call JSON is excluded from
    the loss while surrounding assistant prose stays supervised.
    Returns the modified loss_mask (also mutated in place).
    """
    ids = input_ids.tolist()
    open_ids  = tokenizer.encode("<tool_call>",  add_special_tokens=False)
    close_ids = tokenizer.encode("</tool_call>", add_special_tokens=False)

    cursor = 0
    while True:
        o = _find_subseq(ids, open_ids, cursor)
        if o == -1:
            break
        c = _find_subseq(ids, close_ids, o + len(open_ids))
        if c == -1:
            # Unterminated block: mask to end of sequence to be safe.
            loss_mask[o:] = False
            break
        span_end = c + len(close_ids)          # exclusive
        loss_mask[o:span_end] = False          # delimiters + JSON masked out
        cursor = span_end

    return loss_mask

input_ids, loss_mask = build_chatml_loss_mask(tokenizer, messages)
loss_mask = mask_tool_call_json(input_ids, loss_mask, tokenizer)
labels = labels_from_mask(input_ids, loss_mask)   # -100 at masked positions

6. (Conceptual / short calculation) Checklist item 3 says a quick sanity check before training is (labels != -100).float().mean(), which should land in roughly \(0.25\)-\(0.50\) for chat data. For the 26-token worked example (3 supervised tokens) this quantity is \(3/26 \approx 0.115\) — well below that band. Give two reasons the chapter’s own text explains why a single tiny example falls under the typical range, and describe one dataset-level situation that would push the aggregate fraction above \(0.50\).