fix(sft): mask loss on non-assistant tokens (assistant_only_loss + chat-template markers)

README.md

@@ -244,6 +244,41 @@ Templates convert the list of messages into a single string for the model.
 - **Config**: `data.chat_template: "name"`
 - **Source**: Jinja files located in `src/post_training/chat_templates/templates/`
 
+##### SFT requires `{% generation %}` markers
+
+SFT in this framework uses TRL's `assistant_only_loss=True`, which masks the
+cross-entropy loss on every non-assistant token (system + user). This depends on
+the chat template wrapping the assistant content emission in
+`{% generation %}…{% endgeneration %}` markers — transformers'
+`apply_chat_template(..., return_assistant_tokens_mask=True)` uses them to build
+the per-token loss mask.
+
+If your template lacks those markers, `build_sft_trainer()` will refuse to start
+with a `ValueError` that names the template and points to the fix. This is
+deliberate: silently training without the mask measurably degrades downstream
+performance (the framework previously had this bug; SFT computed CE loss on
+every token in the packed sequence).
+
+Templates that are safe for SFT today:
+
+| Name | Source | Notes |
+|------|--------|-------|
+| `olmo3-instruct-sft` | `allenai/OLMo-3-7B-Instruct-SFT` (HF Hub) | Use to reproduce the Instruct-SFT recipe. |
+| `olmo3-think-sft`    | `allenai/Olmo-3-7B-Think-SFT` (HF Hub)    | Use to reproduce the Think-SFT recipe. |
+
+Templates that are *not* safe for SFT (kept for inference / DPO compatibility):
+
+| Name | Notes |
+|------|-------|
+| `olmo3` | Legacy alias for the markerless Think-SFT template; preserved for inference parity only. |
+| `chatml`, `tulu3`, `apertus` | Markerless; need `{% generation %}` markers added before they can be used for SFT. |
+
+To use a custom template for SFT, wrap exactly the tokens that should contribute
+to loss — typically `content` + `function_calls`/`tool_calls` + the closing
+`<|im_end|>` or `eos_token`. Do **not** wrap the leading role-tag prefix
+(`<|im_start|>assistant\n`); it's a deterministic control sequence the model
+shouldn't have to predict.
+
 #### D. Data inspection
 
 Use the data script to debug the pipeline stages (Raw → Transformed → Formatted → Tokenized) and to compute token statistics.
@@ -431,7 +466,7 @@ checkpointing:
 
 # -- Data mix ----------------------------------------------------------------
 data:
-  chat_template: "olmo3"                     # Name from chat template registry
+  chat_template: "olmo3-instruct-sft"        # Name from chat template registry
   num_proc: null                             # null = auto-detect, capped at 32
   datasets:
     - name: "nemotron_pt_v2"

configs/trl/sft.yaml

@@ -68,7 +68,7 @@ checkpointing:
 
 # -- Data mix ----------------------------------------------------------------
 data:
-  chat_template: "olmo3"                     # Name from chat template registry
+  chat_template: "olmo3-instruct-sft"        # Name from chat template registry
   num_proc: null                             # null = auto-detect, capped at 32
   datasets:
     - name: "nemotron_pt_v2"

src/post_training/chat_templates/registry.py

@@ -10,17 +10,37 @@
 from __future__ import annotations
 
 import logging
+import re
 from pathlib import Path
 
 logger = logging.getLogger(__name__)
 
+# `{% generation %}` markers are a transformers-specific extension to Jinja2,
+# used by ``apply_chat_template(..., return_assistant_tokens_mask=True)`` (and
+# therefore by TRL's ``assistant_only_loss=True``).  Accept any of the four
+# whitespace-stripping variants (``{%``/``{%-`` and ``%}``/``-%}``).
+_GENERATION_OPEN_RE = re.compile(r"\{%-?\s*generation\s*-?%\}")
+_GENERATION_CLOSE_RE = re.compile(r"\{%-?\s*endgeneration\s*-?%\}")
+
 # Directory that holds the .jinja template files.
 _TEMPLATES_DIR = Path(__file__).resolve().parent / "templates"
 
 # Mapping: template name -> jinja filename (relative to _TEMPLATES_DIR).
+#
+# Note on the two ``olmo3*`` entries: they correspond to two *different*
+# AllenAI checkpoints with materially different chat templates.
+#   - ``olmo3``: copied from ``allenai/Olmo-3-7B-Think-SFT``.  Appends
+#     ``<think>`` to ``add_generation_prompt=True`` (Think-style priming).
+#     No ``{% generation %}`` markers, so SFT here cannot mask user/system
+#     tokens out of the loss — the runtime guard in ``methods/sft.py`` will
+#     refuse to start training with this template.
+#   - ``olmo3-instruct-sft``: byte-identical (modulo spliced ``{% generation %}``
+#     markers) to ``allenai/OLMo-3-7B-Instruct-SFT``'s ``chat_template.jinja``.
+#     This is the correct template for reproducing the Instruct-SFT recipe.
 CHAT_TEMPLATES: dict[str, str] = {
     "chatml": "chatml.jinja",
     "olmo3": "olmo3.jinja",
+    "olmo3-instruct-sft": "olmo3-instruct-sft.jinja",
     "apertus": "apertus.jinja",
     "tulu3": "tulu3.jinja",
 }
@@ -41,6 +61,19 @@ def register_chat_template(name: str, filename: str) -> None:
     CHAT_TEMPLATES[name] = filename
 
 
+def has_generation_markers(template: str | None) -> bool:
+    """Return ``True`` if *template* wraps content in
+    ``{% generation %}…{% endgeneration %}`` markers (any whitespace-strip form).
+
+    Required by transformers' ``return_assistant_tokens_mask`` path, which TRL
+    uses to implement ``assistant_only_loss=True``.  Missing markers make the
+    mask silently all-zero — SFT then trains on every token in the sequence.
+    """
+    if not template:
+        return False
+    return bool(_GENERATION_OPEN_RE.search(template) and _GENERATION_CLOSE_RE.search(template))
+
+
 def get_chat_template(name: str) -> str:
     """Return the Jinja source string for the template registered as *name*.
 

src/post_training/chat_templates/templates/olmo3-instruct-sft.jinja

@@ -0,0 +1,16 @@
+{%- set has_system = messages|selectattr('role', 'equalto', 'system')|list|length > 0 -%}{%- if not has_system -%}{{- '<|im_start|>system
+You are a helpful function-calling AI assistant. ' -}}{%- if tools is none or (tools | length) == 0 -%}{{- 'You do not currently have access to any functions. <functions></functions><|im_end|>
+' -}}{%- else -%}{{- 'You are provided with function signatures within <functions></functions> XML tags. You may call one or more functions to assist with the user query. Output any function calls within <function_calls></function_calls> XML tags. Do not make assumptions about what values to plug into functions.' -}}{{- '<functions>' -}}{{- tools | tojson -}}{{- '</functions><|im_end|>
+' -}}{%- endif -%}{%- endif -%}{%- for message in messages -%}{%- if message['role'] == 'system' -%}{{- '<|im_start|>system
+' + message['content'] -}}{%- if tools is not none -%}{{- '<functions>' -}}{{- tools | tojson -}}{{- '</functions>' -}}{%- elif message.get('functions', none) is not none -%}{{- ' <functions>' + message['functions'] + '</functions>' -}}{%- endif -%}{{- '<|im_end|>
+' -}}{%- elif message['role'] == 'user' -%}{{- '<|im_start|>user
+' + message['content'] + '<|im_end|>
+' -}}{%- elif message['role'] == 'assistant' -%}{{- '<|im_start|>assistant
+' -}}{%- generation -%}{%- if message.get('content', none) is not none -%}{{- message['content'] -}}{%- endif -%}{%- if message.get('function_calls', none) is not none -%}{{- '<function_calls>' + message['function_calls'] + '</function_calls>' -}}{% elif message.get('tool_calls', none) is not none %}{{- '<function_calls>' -}}{%- for tool_call in message['tool_calls'] %}{%- if tool_call is mapping and tool_call.get('function', none) is not none %}{%- set args = tool_call['function']['arguments'] -%}{%- set ns = namespace(arguments_list=[]) -%}{%- for key, value in args.items() -%}{%- set ns.arguments_list = ns.arguments_list + [key ~ '=' ~ (value | tojson)] -%}{%- endfor -%}{%- set arguments = ns.arguments_list | join(', ') -%}{{- tool_call['function']['name'] + '(' + arguments + ')' -}}{%- if not loop.last -%}{{ '
+' }}{%- endif -%}{% else %}{{- tool_call -}}{%- endif %}{%- endfor %}{{- '</function_calls>' -}}{%- endif -%}{%- if not loop.last -%}{{- '<|im_end|>' + '
+' -}}{%- else -%}{{- eos_token -}}{%- endif -%}{%- endgeneration -%}{%- elif message['role'] == 'environment' -%}{{- '<|im_start|>environment
+' + message['content'] + '<|im_end|>
+' -}}{%- elif message['role'] == 'tool' -%}{{- '<|im_start|>environment
+' + message['content'] + '<|im_end|>
+' -}}{%- endif -%}{%- if loop.last and add_generation_prompt -%}{{- '<|im_start|>assistant
+' -}}{%- endif -%}{%- endfor -%}

src/post_training/methods/sft.py

@@ -9,6 +9,7 @@
 from accelerate import PartialState
 from trl import SFTConfig, SFTTrainer
 
+from post_training.chat_templates.registry import has_generation_markers
 from post_training.data.loader import load_and_mix_datasets
 from post_training.methods.common import (
     build_callbacks,
@@ -47,6 +48,34 @@ def build_sft_trainer(config: PostTrainingConfig, run_dir: Path) -> SFTTrainer:
     mc = config.sft  # method-specific config
 
     tokenizer = build_tokenizer(config)
+
+    # Fail fast if the chat template can't drive `assistant_only_loss=True`.
+    # Missing markers silently degrade SFT to full-sequence loss — a 21h run
+    # produces a measurably worse model and nothing in the logs shouts.
+    if not has_generation_markers(tokenizer.chat_template):
+        raise ValueError(
+            f"Chat template '{config.data.chat_template}' is missing "
+            "{% generation %}…{% endgeneration %} markers\n"
+            "around the assistant content emission. Without them, "
+            "`assistant_only_loss=True`\n"
+            "is a silent no-op — SFT would compute CE loss on every "
+            "token in the sequence\n"
+            "(system + user + assistant).\n"
+            "\n"
+            "To fix:\n"
+            "  • Switch to a registered marker-bearing template:\n"
+            '      data.chat_template: "olmo3-instruct-sft"   '
+            "# AllenAI OLMo-3-Instruct-SFT recipe\n"
+            '      data.chat_template: "olmo3-think-sft"      '
+            "# AllenAI OLMo-3-Think-SFT recipe\n"
+            "  • Or add `{% generation %}…{% endgeneration %}` markers "
+            "around the assistant\n"
+            "    content emission in your own jinja template.\n"
+            "\n"
+            "Reference: open-instruct's sft_tulu_tokenize_and_truncate_v1\n"
+            "(open-instruct/open_instruct/dataset_transformation.py L1111-L1176)."
+        )
+
     with PartialState().main_process_first():
         dataset = load_and_mix_datasets(config.data, row_filter=_sft_row_filter)
 
@@ -56,6 +85,10 @@ def build_sft_trainer(config: PostTrainingConfig, run_dir: Path) -> SFTTrainer:
         packing=mc.packing,
         dataset_num_proc=mc.dataset_num_proc,
         model_init_kwargs=build_model_init_kwargs(config),
+        # Mask loss on everything except the assistant content.  Requires the
+        # chat template to wrap assistant turns in {% generation %}…{% endgeneration %}.
+        # Without this, SFTTrainer trains on user + system tokens too.
+        assistant_only_loss=True,
     )
 
     trainer = SFTTrainer(