ChatDataset answer-only loss mask starts 3 tokens too early without {% generation %}, teaching the model to re-emit the turn header

Summary

When a chat template does not contain {% generation %} markers, format_chat_template falls back to _build_multiturn_assistant_mask, which locates the supervised region by prefix-length arithmetic. That arithmetic computes the length of the rendered prefix without add_generation_prompt=True, so the supervised region starts at the beginning of the generation-prompt header instead of after it.

The model is therefore trained to emit the turn header as the first thing in its response. This affects the documented DiffusionGemma SFT/LoRA guide, which uses google/diffusiongemma-26B-A4B-it — a model whose header contains the ordinary (non-special) token model, so the artifact survives detokenization and is visible in generated text.

Reproduction

Uses only public artifacts: one openai/gsm8k row (the schema examples/dllm_sft/prep_gsm8k.py produces), the official tokenizer and chat template, and _build_multiturn_assistant_mask itself.

from transformers import AutoTokenizer
from nemo_automodel.components.datasets.llm.formatting_utils import (
    GENERATION_REGEX, _build_multiturn_assistant_mask,
)

tok = AutoTokenizer.from_pretrained("google/diffusiongemma-26B-A4B-it")

# The exact schema examples/dllm_sft/prep_gsm8k.py writes
messages = [
    {"role": "user", "content": "Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?"},
    {"role": "assistant", "content": "Natalia sold 48/2 = <<48/2=24>>24 clips in May.\nNatalia sold 48+24 = <<48+24=72>>72 clips altogether in April and May.\n#### 72"},
]

print("template has generation markers:", GENERATION_REGEX.search(tok.chat_template) is not None)

full_ids = tok.apply_chat_template(messages, tokenize=True)
mask = _build_multiturn_assistant_mask(tok, messages, full_ids)
mask_start = mask.index(1)

# Where the model actually takes over at inference time
infer_len = len(tok.apply_chat_template(messages[:1], tokenize=True, add_generation_prompt=True))

print("mask start :", mask_start)
print("prompt len :", infer_len)
print("wrongly supervised:", [(i, tok.decode([full_ids[i]])) for i in range(mask_start, infer_len)])
assert mask_start == infer_len, f"off by {infer_len - mask_start} tokens"

Output:

template has generation markers: False
mask start : 43
prompt len : 46
wrongly supervised: [(43, '<|turn>'), (44, 'model'), (45, '\n')]
AssertionError: off by 3 tokens

Root cause

_tokenize_chat has no add_generation_prompt parameter, so prefix_length(k) measures "length of the conversation rendered through message k" rather than "length of the prompt as the model sees it at inference". The two relevant places in nemo_automodel/components/datasets/llm/formatting_utils.py:

• _tokenize_chat — calls apply_chat_template without add_generation_prompt
• _build_multiturn_assistant_mask — start = prefix_length(idx)

These two quantities differ by exactly the length of the generation prompt, which every chat template emits to mark turn ownership. The assumption "message boundary == generation boundary" does not hold for any realistic template; only the size of the gap is model-specific (3 tokens for DiffusionGemma: <|turn>, model, \n).

Why the existing guard does not catch it

_is_consistent_render_prefix validates that a rendered prefix is a genuine prefix of the full render. That catches templates which rewrite earlier turns based on later ones (the Qwen3 <think>-dropping case cited in the error message), but it does not validate the position of the boundary. DiffusionGemma's template never rewrites history, so the check passes while the boundary is still off by 3.

Notably the error message raised by that guard already recommends the correct fix ("Provide a chat template that wraps assistant turns in {% generation %}"), but nothing enforces it, and the documented example path silently takes the unsafe branch.

Impact

Any ChatDataset run whose template lacks {% generation %} supervises the generation-prompt header. The consequences are easy to miss:

• The extra tokens are constant and trivially learnable, so they lower average loss slightly rather than showing up as a regression. Loss curves look fine.
• If the header consists entirely of special tokens, detokenization hides the artifact.
• GSM8K accuracy is graded by extracting the value after ####, so a leading model\n does not change the score.

The documented DiffusionGemma guide reports only loss curves for the first 200 steps and never inspects a generation, which is why this has not surfaced there. On a long-form generation task scored by n-gram overlap the leaked prefix is immediately visible and costs metric points.

Suggested fixes

1. Assert the boundary in the fallback path. For each assistant turn, require start == len(render(messages[:idx], add_generation_prompt=True)) and raise the existing "use {% generation %}" error when it does not hold. This turns a silent data-quality bug into a startup failure.
2. Alternatively, use add_generation_prompt=True when computing the start offset, so the arithmetic measures the real handover point.
3. Ship a {% generation %} template for the DiffusionGemma examples (or document the requirement in docs/guides/dllm/diffusiongemma.mdx), so the documented path uses the exact assistant_masks branch.

I have a working template patch for google/diffusiongemma-26B-A4B-it that wraps the assistant content and its turn-close in {% generation %} and renders byte-identically to the official template. I'll open a PR alongside this issue.

Environment

• transformers 4.57.6
• tokenizer/chat_template: google/diffusiongemma-26B-A4B-it @ main
• affected config examples: examples/dllm_sft/diffusion_gemma_sft.yaml, examples/dllm_sft/diffusion_gemma_lora.yaml (both set dataset.tokenizer.pretrained_model_name_or_path: google/diffusiongemma-26B-A4B-it and rely on ChatDataset + mask_history: true)

Start by running the reproduction with _build_multiturn_assistant_mask and the DiffusionGemma tokenizer. Inspect _tokenize_chat and _build_multiturn_assistant_mask in nemo_automodel/components/datasets/llm/formatting_utils.py, then review the DiffusionGemma SFT and LoRA configs and docs. Done means the fallback boundary matches the inference prompt or fails safely, with regression coverage for the reported three-token offset.