Chat templates, loss masks and response contracts

A fine-tuning row is not just a question and an answer. The model sees a serialized conversation with role markers, separators and termination tokens. Those bytes determine where generation starts and which parts of the conversation the loss teaches it to reproduce.

I treat the tokenizer, chat template, loss mask and response parser as one contract. Training and serving must agree on the conversation structure. The safest way to establish that agreement is to inspect rendered examples and test their token boundaries before launching a long run.

Render the conversation in two forms

Start with the same user message twice. The first rendering ends at the assistant-generation prefix. The second includes the completed assistant response. Tokenize both without adding another layer of special tokens. The first sequence must be an exact prefix of the second for the simple completion-only masking method below.

def encode_completion(processor, prompt, completion, max_length=2048):
    tokenizer = processor.tokenizer
    user = [{"role": "user", "content": prompt}]
    full = user + [{"role": "assistant", "content": completion}]
    prefix_text = processor.apply_chat_template(
        user, tokenize=False, add_generation_prompt=True,
        enable_thinking=False,
    )
    full_text = processor.apply_chat_template(
        full, tokenize=False, add_generation_prompt=False,
        enable_thinking=False,
    )
    prefix = tokenizer(prefix_text, add_special_tokens=False)["input_ids"]
    ids = tokenizer(full_text, add_special_tokens=False)["input_ids"]
    if not prefix or ids[:len(prefix)] != prefix:
        raise ValueError("template prefix mismatch")
    if len(ids) > max_length or len(ids) == len(prefix):
        raise ValueError("invalid supervision length")
    return {
        "input_ids": ids,
        "attention_mask": [1] * len(ids),
        "labels": [-100] * len(prefix) + ids[len(prefix):],
    }

The prefix assertion is essential. A template may render an assistant prefix differently when the final answer is present. If that happens, use a template-aware masking mechanism or an explicitly verified serialization; do not guess the boundary from a token count.

Transformers’ Gemma 4 documentation identifies the processor for this architecture. The artifact shipped to serving must carry the matching tokenizer and template, not merely weights that happen to load.

See the labels directly

Here is the masking operation using tiny token IDs so its behaviour is obvious. These IDs are a teaching example, not Gemma vocabulary entries.

from corpus import completion_labels

prompt_ids = [11, 12, 13]
full_ids = [11, 12, 13, 21, 22, 2]
labels = completion_labels(prompt_ids, full_ids, max_length=8)
assert labels == [-100, -100, -100, 21, 22, 2]

The model still attends to prompt tokens. Ignoring them in the loss does not remove them from context. It means that the optimizer is rewarded for predicting the completion rather than for reproducing the source document.

Training frameworks perform the causal shift internally. Do not manually shift labels again unless the model implementation explicitly requires it. A second shift changes the prediction target and can produce misleading loss curves.

Thinking mode is a model-format decision

A model’s native thinking mode has its own control tokens and template behaviour. Use the processor’s supported control rather than inventing a tag and assuming it activates the same mechanism. The reference adaptation recipe sets enable_thinking=False to keep its response protocol bounded. The Gemma thinking guide describes the native control.

If a separate experiment enables thinking, record that change as part of the candidate configuration. Account for all generated tokens in latency and cost. A longer internal generation phase may affect the time to a useful final answer even if the UI displays only the final response.

Training explanations are also distinct from a native reasoning channel. A concise explanation such as “₹30 lakh divided by ₹50 lakh is 60%” is a supervised answer artifact. It should not be presented as privileged access to the model’s internal reasoning process.

Choose an output contract the runtime can execute

The reference runner uses a small application-level JSON protocol. A decision is either a tool request or a final result. It is deliberately separate from provider-native function-calling formats, so the same sandbox can exercise different local models through a common adapter.

{
  "kind": "tool",
  "name": "future_value",
  "arguments": {
    "principal_inr": "2500000",
    "annual_rate": "0.06",
    "years": 5
  }
}

The final result contains an outcome, typed claims, evidence IDs and a policy code. The runtime supplies the actual tool trace separately. It never accepts a model-generated assertion that a tool executed successfully.

A production UI can render the final structured object through fixed templates or add a separately validated natural-language explanation. This reference protocol evaluates structured outcomes. It does not by itself evaluate whether a Hindi paragraph sounds natural; that requires an additional language-quality evaluation.

Make multi-turn supervision deliberate

A trajectory contains a user request, an assistant tool decision, a tool observation and an assistant answer. For multi-turn fine-tuning, choose whether to supervise all assistant spans or only the final completion. Tool observations usually remain inputs, not generated targets.

The supplied training entry point implements a prompt/completion format. To teach a final response after a tool result, serialize the reviewed trajectory context into the prompt and supervise the final response. For native multi-turn training, extend the renderer with explicit assistant masks and test every span.

TRL’s SFT documentation distinguishes completion-only loss from assistant-only masking and describes the template support needed for assistant masks. Selecting the option is not enough; inspect the resulting labels on an actual batch.

Packing multiple examples introduces another boundary. Without the intended attention separation, one example can see tokens from another. The reference recipe leaves packing off. If you enable it for throughput, document the packing strategy and test that example boundaries preserve the intended learning task.

Eight checks before an optimizer step

For a sampled batch, decode the supervised tokens and read them. This is one of the rare debugging steps where a few minutes of human inspection can prevent an entire wasted training run.

Record the proportion of supervised tokens as well. If the average sequence is 1,800 tokens but only 80 are supervised, “tokens processed” and “answer tokens learned” describe different quantities. Compare runs using both.

Preserve the contract through export

After adapter selection, render a fixed set of conversations in the training environment and the serving environment. Compare the intended roles, template settings, generation prefix and stopping behaviour. Byte-for-byte equality is ideal when the same tokenizer implementation is used; when a serving conversion changes representation, establish equivalent token boundaries with explicit probes.

Then run behavioural fixtures. The exported model must still produce valid decisions, respect the output budget and stop at the expected boundary. A correct adapter attached to the wrong base or template is a different candidate.

The release manifest therefore includes template and tokenizer digests, thinking mode, maximum output tokens, stop configuration and parser version. These small artifacts deserve the same versioning discipline as the large weight files. They directly determine whether the model’s output becomes an actionable, reviewable result.