feat: add Qwen3-8B template

configs/trl/sft.yaml

@@ -65,6 +65,11 @@ training:
 sft:
   max_seq_length: 4096
   packing: true
+  truncated_span_action: warn                # warn | drop — a row whose assistant span
+                                             # starts inside max_seq_length and runs past
+                                             # it trains on an answer cut mid-sentence.
+                                             # "drop" removes those rows, for when the
+                                             # sequence length is fixed by memory.
 
 # -- Checkpointing -----------------------------------------------------------
 checkpointing:

scripts/data.py

@@ -26,6 +26,7 @@
 import argparse
 import statistics
 import sys
+from functools import partial
 from pathlib import Path
 from typing import Any
 
@@ -256,7 +257,7 @@ def _run_count_tokens(args: argparse.Namespace, cli_overrides: list[str]) -> Non
     # ── Lazy imports (heavy) ────────────────────────────────────────
     from post_training.data.loader import _resolve_num_proc, load_and_mix_datasets
     from post_training.methods.common import build_tokenizer
-    from post_training.methods.sft import MESSAGES_FEATURES, _sft_row_filter
+    from post_training.methods.sft import MESSAGES_FEATURES, _filter_sft_rows
 
     # Load config
     config = PostTrainingConfig.load(args.config, cli_overrides)
@@ -267,16 +268,20 @@ def _run_count_tokens(args: argparse.Namespace, cli_overrides: list[str]) -> Non
 
     columns_to_keep = None
     features = None
-    row_filter = None
+    dataset_filter_fn = None
     if config.method == "sft":
         columns_to_keep = ["messages"]
         features = MESSAGES_FEATURES
-        row_filter = _sft_row_filter
+        dataset_filter_fn = partial(
+            _filter_sft_rows,
+            tokenizer=tokenizer,
+            max_length=config.sft.max_seq_length,
+        )
 
     # Load dataset mix
     ds = load_and_mix_datasets(
         config.data,
-        row_filter=row_filter,
+        dataset_filter_fn=dataset_filter_fn,
         columns_to_keep=columns_to_keep,
         features=features,
     )

src/post_training/chat_templates/registry.py

@@ -46,13 +46,23 @@
 #     recipe via TRL.  Same assistant-only masking pattern as Instruct-SFT;
 #     OLMo-core's reference Think pipeline uses identical offline-baked
 #     masks (only learning rate and dataset differ between the two recipes).
+#
+# Note on ``qwen3``: Unlike the OLMo templates, ``qwen3`` masks *only the
+# final exchange*: Qwen3 strips ``<think>`` blocks from assistant turns at
+# or before the last user query, so history renders as a bare answer with
+# no reasoning trace. Training on that teaches the model to answer directly at
+# exactly the position where it opens a ``<think>`` block at inference.  The
+# markers therefore wrap the assistant body only when
+# ``loop.index0 > ns.last_query_index``.  Multi-step tool exchanges keep every
+# assistant turn after the last user query, since those retain their traces.
 CHAT_TEMPLATES: dict[str, str] = {
     "chatml": "chatml.jinja",
     "olmo3": "olmo3.jinja",
     "olmo3-instruct-sft": "olmo3-instruct-sft.jinja",
     "olmo3-think-sft": "olmo3-think-sft.jinja",
     "apertus": "apertus.jinja",
     "tulu3": "tulu3.jinja",
+    "qwen3": "qwen3.jinja",
 }
 
 

src/post_training/chat_templates/templates/qwen3.jinja

@@ -0,0 +1,98 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0].role == 'system' %}
+        {{- messages[0].content + '\n\n' }}
+    {%- endif %}
+    {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
+        {%- set ns.multi_step_tool = false %}
+        {%- set ns.last_query_index = index %}
+    {%- endif %}
+{%- endfor %}
+{%- for message in messages %}
+    {%- if message.content is string %}
+        {%- set content = message.content %}
+    {%- else %}
+        {%- set content = '' %}
+    {%- endif %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- set assistant_body %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {%- if loop.last or (not loop.last and reasoning_content) %}
+                {{- '<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
+            {%- else %}
+                {{- content }}
+            {%- endif %}
+        {%- else %}
+            {{- content }}
+        {%- endif %}
+        {%- if message.tool_calls %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if (loop.first and content) or (not loop.first) %}
+                    {{- '\n' }}
+                {%- endif %}
+                {%- if tool_call.function %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {{- '<tool_call>\n{"name": "' }}
+                {{- tool_call.name }}
+                {{- '", "arguments": ' }}
+                {%- if tool_call.arguments is string %}
+                    {{- tool_call.arguments }}
+                {%- else %}
+                    {{- tool_call.arguments | tojson }}
+                {%- endif %}
+                {{- '}\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>' }}
+        {%- endset %}
+        {{- '<|im_start|>' + message.role + '\n' }}
+        {%- if loop.index0 > ns.last_query_index %}
+            {%- generation %}{{- assistant_body -}}{%- endgeneration %}
+        {%- else %}
+            {{- assistant_body }}
+        {%- endif %}
+        {{- '\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- endif %}
+{%- endif %}

src/post_training/config.py

@@ -117,6 +117,13 @@ class SFTMethodConfig:
     max_seq_length: int = 4096
     packing: bool = True
     dataset_num_proc: int | None = None
+    # What to do with a row whose supervised span starts inside max_seq_length
+    # and continues past it. Such a row trains on an answer cut mid-sentence.
+    # "warn" keeps it and says so; "drop" removes it, for when the sequence
+    # length is fixed by memory and shortening the data is the only option.
+    # Dropping is never silent: it changes the realised mixture, because
+    # data.datasets[].weight applies to the rows that survive filtering.
+    truncated_span_action: str = "warn"
 
 
 @dataclass

src/post_training/data/loader.py

@@ -101,7 +101,7 @@ def _resample_to_size(ds: Dataset, target_n: int, seed: int) -> Dataset:
 
 def load_and_mix_datasets(
     config: DataConfig,
-    row_filter: Callable[[dict], bool] | None = None,
+    dataset_filter_fn: Callable[[Dataset, int], Dataset] | None = None,
     columns_to_keep: list[str] | None = None,
     features: Features | None = None,
 ) -> Dataset:
@@ -111,10 +111,12 @@ def load_and_mix_datasets(
     ----------
     config:
         The ``data`` section of :class:`PostTrainingConfig`.
-    row_filter:
-        Optional predicate applied after transforms to exclude invalid
-        rows.  Each training method passes its own filter (e.g. SFT
-        checks ``messages``, DPO checks ``chosen`` / ``rejected``).
+    dataset_filter_fn:
+        Optional callable ``(dataset, num_proc) -> dataset`` applied to each
+        dataset after transforms and column selection. Each training method
+        passes its own filter (e.g. SFT drops rows whose assistant tokens
+        never enter the loss, DPO drops rows with an empty ``chosen`` or
+        ``rejected``). Bind extra arguments with :func:`functools.partial`.
         When ``None``, no filtering is applied.
     columns_to_keep:
         Optional list of columns to retain after loading/mapping. When a
@@ -205,10 +207,12 @@ def load_and_mix_datasets(
                 )
             ds = ds.select_columns(present)
 
-        # Apply method-specific row filter (e.g. SFT checks for non-empty
-        # "messages", DPO checks for non-empty "chosen" / "rejected").
-        if row_filter is not None:
-            ds = ds.filter(row_filter, num_proc=num_proc)
+        # Apply the method-specific dataset filter (e.g. SFT drops rows with
+        # an all-zero assistant loss mask, DPO drops empty preference pairs).
+        if dataset_filter_fn is not None:
+            logger.info("Filtering dataset '%s' (%d rows).", entry.name, len(ds))
+            ds = dataset_filter_fn(ds, num_proc=num_proc)
+            logger.info("Dataset '%s' has %d rows after filtering.", entry.name, len(ds))
 
         loaded_datasets.append(ds)
 

src/post_training/methods/dpo.py

@@ -7,6 +7,7 @@
 from typing import TYPE_CHECKING
 
 from accelerate import PartialState
+from datasets import Dataset
 from trl import DPOConfig, DPOTrainer
 
 from post_training.data.loader import load_and_mix_datasets
@@ -25,9 +26,13 @@
 logger = logging.getLogger(__name__)
 
 
-def _dpo_row_filter(example: dict) -> bool:
-    """Keep only rows with non-empty ``chosen`` and ``rejected`` fields."""
-    return len(example.get("chosen", [])) > 0 and len(example.get("rejected", [])) > 0
+def _filter_dpo_rows(ds: Dataset, num_proc: int) -> Dataset:
+    """Drop rows with an empty ``chosen`` or ``rejected`` field."""
+    return ds.filter(
+        lambda row: len(row["chosen"]) > 0 and len(row["rejected"]) > 0,
+        num_proc=num_proc,
+        desc="filtering empty preference pairs",
+    )
 
 
 def build_dpo_trainer(config: PostTrainingConfig, run_dir: Path) -> DPOTrainer:
@@ -49,7 +54,7 @@ def build_dpo_trainer(config: PostTrainingConfig, run_dir: Path) -> DPOTrainer:
 
     tokenizer = build_tokenizer(config)
     with PartialState().main_process_first():
-        dataset = load_and_mix_datasets(config.data, row_filter=_dpo_row_filter)
+        dataset = load_and_mix_datasets(config.data, dataset_filter_fn=_filter_dpo_rows)
 
     dpo_config = DPOConfig(
         **build_common_training_kwargs(config, run_dir),