From 82126f619872c883459e5c0ea737a0ad4024e120 Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Thu, 6 Aug 2026 17:04:34 +0200 Subject: [PATCH 01/12] feat: add original Qwen3-8B template from HuggingFace --- .../chat_templates/templates/qwen3-8b.jinja | 89 +++++++++++++++++++ 1 file changed, 89 insertions(+) create mode 100644 src/post_training/chat_templates/templates/qwen3-8b.jinja diff --git a/src/post_training/chat_templates/templates/qwen3-8b.jinja b/src/post_training/chat_templates/templates/qwen3-8b.jinja new file mode 100644 index 0000000..01be9b3 --- /dev/null +++ b/src/post_training/chat_templates/templates/qwen3-8b.jinja @@ -0,0 +1,89 @@ +{%- 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 XML tags:\n" }} + {%- for tool in tools %} + {{- "\n" }} + {{- tool | tojson }} + {%- endfor %} + {{- "\n\n\nFor each function call, return a json object with function name and arguments within XML tags:\n\n{\"name\": , \"arguments\": }\n<|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('') and message.content.endswith('')) %} + {%- 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 '' in content %} + {%- set reasoning_content = content.split('')[0].rstrip('\n').split('')[-1].lstrip('\n') %} + {%- set content = content.split('')[-1].lstrip('\n') %} + {%- endif %} + {%- endif %} + {%- if loop.index0 > ns.last_query_index %} + {%- if loop.last or (not loop.last and reasoning_content) %} + {{- '<|im_start|>' + message.role + '\n\n' + reasoning_content.strip('\n') + '\n\n\n' + content.lstrip('\n') }} + {%- else %} + {{- '<|im_start|>' + message.role + '\n' + content }} + {%- endif %} + {%- else %} + {{- '<|im_start|>' + message.role + '\n' + 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 %} + {{- '\n{"name": "' }} + {{- tool_call.name }} + {{- '", "arguments": ' }} + {%- if tool_call.arguments is string %} + {{- tool_call.arguments }} + {%- else %} + {{- tool_call.arguments | tojson }} + {%- endif %} + {{- '}\n' }} + {%- endfor %} + {%- endif %} + {{- '<|im_end|>\n' }} + {%- elif message.role == "tool" %} + {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %} + {{- '<|im_start|>user' }} + {%- endif %} + {{- '\n\n' }} + {{- content }} + {{- '\n' }} + {%- 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 %} + {{- '\n\n\n\n' }} + {%- endif %} +{%- endif %} \ No newline at end of file From 3fa3ee1f209ec08d3ab2866aa1754ff3b0478d12 Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Thu, 6 Aug 2026 17:18:37 +0200 Subject: [PATCH 02/12] feat: add generation markers --- .../chat_templates/templates/qwen3-8b.jinja | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/src/post_training/chat_templates/templates/qwen3-8b.jinja b/src/post_training/chat_templates/templates/qwen3-8b.jinja index 01be9b3..a435a14 100644 --- a/src/post_training/chat_templates/templates/qwen3-8b.jinja +++ b/src/post_training/chat_templates/templates/qwen3-8b.jinja @@ -40,14 +40,16 @@ {%- set content = content.split('')[-1].lstrip('\n') %} {%- endif %} {%- endif %} + {{- '<|im_start|>' + message.role + '\n' }} + {%- generation %} {%- if loop.index0 > ns.last_query_index %} {%- if loop.last or (not loop.last and reasoning_content) %} - {{- '<|im_start|>' + message.role + '\n\n' + reasoning_content.strip('\n') + '\n\n\n' + content.lstrip('\n') }} + {{- '\n' + reasoning_content.strip('\n') + '\n\n\n' + content.lstrip('\n') }} {%- else %} - {{- '<|im_start|>' + message.role + '\n' + content }} + {{- content }} {%- endif %} {%- else %} - {{- '<|im_start|>' + message.role + '\n' + content }} + {{- content }} {%- endif %} {%- if message.tool_calls %} {%- for tool_call in message.tool_calls %} @@ -68,7 +70,9 @@ {{- '}\n' }} {%- endfor %} {%- endif %} - {{- '<|im_end|>\n' }} + {{- '<|im_end|>' }} + {%- endgeneration %} + {{- '\n' }} {%- elif message.role == "tool" %} {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %} {{- '<|im_start|>user' }} From f48454b00405203e1a947b6b47d916f7d0f839d8 Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Thu, 6 Aug 2026 17:25:13 +0200 Subject: [PATCH 03/12] feat: add Qwen3-8B template to registry --- src/post_training/chat_templates/registry.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/post_training/chat_templates/registry.py b/src/post_training/chat_templates/registry.py index 84b0700..4c0d972 100644 --- a/src/post_training/chat_templates/registry.py +++ b/src/post_training/chat_templates/registry.py @@ -53,6 +53,7 @@ "olmo3-think-sft": "olmo3-think-sft.jinja", "apertus": "apertus.jinja", "tulu3": "tulu3.jinja", + "qwen3": "qwen3-8b.jinja", } From 6884b66f28d1e5c89b32d403bfe499d9904e94b7 Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Fri, 7 Aug 2026 13:15:24 +0200 Subject: [PATCH 04/12] fix: exclude think-stripped history from Qwen3 loss mask Qwen3 drops from assistant turns at or before the last user query, so history rendered without a reasoning trace was being trained on. --- src/post_training/chat_templates/registry.py | 16 ++++++++++++++++ .../chat_templates/templates/qwen3-8b.jinja | 11 ++++++++--- 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/src/post_training/chat_templates/registry.py b/src/post_training/chat_templates/registry.py index 4c0d972..3afad4e 100644 --- a/src/post_training/chat_templates/registry.py +++ b/src/post_training/chat_templates/registry.py @@ -46,6 +46,22 @@ # 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``: byte-identical (modulo spliced ``{% generation %}`` +# markers) to the ``chat_template`` field in ``Qwen/Qwen3-8B``'s +# ``tokenizer_config.json``. The same template ships with every model in the +# original Qwen3 release — 0.6B through 235B-A22B, dense and MoE — so the +# entry is not size-specific. The ``*-Instruct-2507`` and ``*-Thinking-2507`` +# refreshes do carry different templates and would need their own entries. +# +# Unlike the OLMo templates, ``qwen3`` masks *only the final exchange*: Qwen3 +# strips ```` 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 ```` 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", diff --git a/src/post_training/chat_templates/templates/qwen3-8b.jinja b/src/post_training/chat_templates/templates/qwen3-8b.jinja index a435a14..a8ff3aa 100644 --- a/src/post_training/chat_templates/templates/qwen3-8b.jinja +++ b/src/post_training/chat_templates/templates/qwen3-8b.jinja @@ -40,8 +40,7 @@ {%- set content = content.split('')[-1].lstrip('\n') %} {%- endif %} {%- endif %} - {{- '<|im_start|>' + message.role + '\n' }} - {%- generation %} + {%- set assistant_body %} {%- if loop.index0 > ns.last_query_index %} {%- if loop.last or (not loop.last and reasoning_content) %} {{- '\n' + reasoning_content.strip('\n') + '\n\n\n' + content.lstrip('\n') }} @@ -71,7 +70,13 @@ {%- endfor %} {%- endif %} {{- '<|im_end|>' }} - {%- endgeneration %} + {%- 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") %} From 877c600cc595a4352382bddedf29a0d599e20796 Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Fri, 7 Aug 2026 13:15:25 +0200 Subject: [PATCH 05/12] test: assert which spans the qwen3 template puts in the loss --- tests/test_chat_template_masking.py | 123 ++++++++++++++++++++++++++++ 1 file changed, 123 insertions(+) diff --git a/tests/test_chat_template_masking.py b/tests/test_chat_template_masking.py index 90c6252..a54d42f 100644 --- a/tests/test_chat_template_masking.py +++ b/tests/test_chat_template_masking.py @@ -7,10 +7,17 @@ We test the detector itself (cheap, no deps) plus per-template regression guards on the registered chat templates. + +Marker *presence* is necessary but not sufficient: markers spliced around +the wrong span mask the wrong tokens while still passing the detector. +The span-level tests at the bottom render templates through transformers' +own tracker and assert on which text actually lands in the loss. """ from __future__ import annotations +import pytest + from post_training.chat_templates.registry import ( get_chat_template, has_generation_markers, @@ -74,3 +81,119 @@ def test_olmo3_template_lacks_markers() -> None: which checkpoint you are reproducing. """ assert not has_generation_markers(get_chat_template("olmo3")) + + +def test_qwen3_template_has_markers() -> None: + """The qwen3 template drives `assistant_only_loss=True` for Qwen3 SFT + runs. Same story as the OLMo templates — no markers, no masking. + """ + assert has_generation_markers(get_chat_template("qwen3")) + + +# ── span-level masking tests ─────────────────────────────────────────── +# +# These render through transformers' generation tracker — the same code +# path `return_assistant_tokens_mask=True` (and therefore TRL's +# `assistant_only_loss=True`) uses — and assert on the exact substrings +# that end up in the loss. Character spans, not token spans, so no +# tokenizer download is needed. + + +def _assistant_spans(template_name: str, messages: list[dict]) -> list[str]: + """Return the substrings of the rendered conversation that are in the loss.""" + # transformers-internal: if this import moves, the public equivalent is + # `tokenizer.apply_chat_template(..., return_assistant_tokens_mask=True)`, + # which additionally requires a tokenizer. + chat_template_utils = pytest.importorskip("transformers.utils.chat_template_utils") + + compiled = chat_template_utils._compile_jinja_template(get_chat_template(template_name)) + rendered, indices = chat_template_utils._render_with_assistant_indices( + compiled, messages, None, None, False + ) + return [rendered[start:end] for start, end in indices] + + +def _think(reasoning: str, answer: str) -> str: + return f"\n{reasoning}\n\n\n{answer}" + + +def test_qwen3_excludes_history_assistant_turns() -> None: + """Qwen3 strips `` from assistant turns at or before the last + user query, so those turns render as a bare answer with no reasoning + trace. Training on them teaches the model to answer directly at + exactly the position where it always opens a `` block at + inference — a train/inference mismatch that trains against the + reasoning behaviour. Only the final exchange belongs in the loss. + """ + messages = [ + {"role": "user", "content": "Q1"}, + {"role": "assistant", "content": _think("R1", "A1")}, + {"role": "user", "content": "Q2"}, + {"role": "assistant", "content": _think("R2", "A2")}, + ] + + assert _assistant_spans("qwen3", messages) == [_think("R2", "A2") + "<|im_end|>"] + + +def test_qwen3_keeps_every_assistant_turn_after_last_user_query() -> None: + """`ns.last_query_index` tracks the last *real* user message, not the + last message. A multi-step tool exchange therefore has several + assistant turns after it, and all of them keep their reasoning trace + when rendered — so all of them belong in the loss. + """ + messages = [ + {"role": "user", "content": "weather?"}, + { + "role": "assistant", + "content": "\nCall the tool.\n\n\n", + "tool_calls": [ + { + "type": "function", + "function": {"name": "get_weather", "arguments": {"city": "Berlin"}}, + } + ], + }, + {"role": "tool", "content": '{"temp_c": 18}'}, + {"role": "assistant", "content": "18C in Berlin."}, + ] + + spans = _assistant_spans("qwen3", messages) + assert len(spans) == 2 + assert spans[0].startswith("\nCall the tool.\n") + assert "" in spans[0] + assert spans[1].endswith("18C in Berlin.<|im_end|>") + + +def test_qwen3_excludes_prompt_header_and_turn_separator() -> None: + """The `<|im_start|>assistant\\n` header is prompt, and the `\\n` after + `<|im_end|>` is a turn separator — neither is something the model + generates, so both stay outside the span. + """ + messages = [ + {"role": "user", "content": "Q"}, + {"role": "assistant", "content": _think("R", "A")}, + ] + + (span,) = _assistant_spans("qwen3", messages) + assert not span.startswith("<|im_start|>") + assert span.endswith("<|im_end|>") + + +def test_olmo3_instruct_sft_masks_all_assistant_turns() -> None: + """Counterpoint to the qwen3 behaviour above, pinned deliberately: the + OLMo templates do not strip anything from history, so every assistant + turn renders exactly as generated and every one belongs in the loss. + The qwen3 exclusion is a Qwen-specific consequence of think-stripping, + not a house rule about multi-turn data. + """ + messages = [ + {"role": "user", "content": "Q1"}, + {"role": "assistant", "content": "A1"}, + {"role": "user", "content": "Q2"}, + {"role": "assistant", "content": "A2"}, + ] + + spans = _assistant_spans("olmo3-instruct-sft", messages) + assert len(spans) == 2 + assert spans[0].startswith("A1") + assert spans[1].startswith("A2") From ec74b1d562d57efdde735fe6eeb12539819ecf41 Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Fri, 7 Aug 2026 13:15:24 +0200 Subject: [PATCH 06/12] fix: exclude think-stripped history from Qwen3 loss mask Qwen3 drops from assistant turns at or before the last user query, so history rendered without a reasoning trace was being trained on. --- src/post_training/chat_templates/registry.py | 9 +++++++++ .../chat_templates/templates/qwen3-8b.jinja | 11 ++++++++--- 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/src/post_training/chat_templates/registry.py b/src/post_training/chat_templates/registry.py index 4c0d972..2a8ddc5 100644 --- a/src/post_training/chat_templates/registry.py +++ b/src/post_training/chat_templates/registry.py @@ -46,6 +46,15 @@ # 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 ```` 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 ```` 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", diff --git a/src/post_training/chat_templates/templates/qwen3-8b.jinja b/src/post_training/chat_templates/templates/qwen3-8b.jinja index a435a14..a8ff3aa 100644 --- a/src/post_training/chat_templates/templates/qwen3-8b.jinja +++ b/src/post_training/chat_templates/templates/qwen3-8b.jinja @@ -40,8 +40,7 @@ {%- set content = content.split('')[-1].lstrip('\n') %} {%- endif %} {%- endif %} - {{- '<|im_start|>' + message.role + '\n' }} - {%- generation %} + {%- set assistant_body %} {%- if loop.index0 > ns.last_query_index %} {%- if loop.last or (not loop.last and reasoning_content) %} {{- '\n' + reasoning_content.strip('\n') + '\n\n\n' + content.lstrip('\n') }} @@ -71,7 +70,13 @@ {%- endfor %} {%- endif %} {{- '<|im_end|>' }} - {%- endgeneration %} + {%- 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") %} From 8d4e67678886576f74ea72557a64d26a38f98d98 Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Fri, 7 Aug 2026 13:15:25 +0200 Subject: [PATCH 07/12] test: assert which spans the qwen3 template puts in the loss --- tests/test_chat_template_masking.py | 125 ++++++++++++++++++++++++++++ 1 file changed, 125 insertions(+) diff --git a/tests/test_chat_template_masking.py b/tests/test_chat_template_masking.py index 90c6252..0c9f7f6 100644 --- a/tests/test_chat_template_masking.py +++ b/tests/test_chat_template_masking.py @@ -7,10 +7,17 @@ We test the detector itself (cheap, no deps) plus per-template regression guards on the registered chat templates. + +Marker *presence* is necessary but not sufficient: markers spliced around +the wrong span mask the wrong tokens while still passing the detector. +The span-level tests at the bottom render templates through transformers' +own tracker and assert on which text actually lands in the loss. """ from __future__ import annotations +import pytest + from post_training.chat_templates.registry import ( get_chat_template, has_generation_markers, @@ -74,3 +81,121 @@ def test_olmo3_template_lacks_markers() -> None: which checkpoint you are reproducing. """ assert not has_generation_markers(get_chat_template("olmo3")) + + +def test_qwen3_template_has_markers() -> None: + """The qwen3 template drives `assistant_only_loss=True` for Qwen3 SFT + runs. Same story as the OLMo templates — no markers, no masking. + """ + assert has_generation_markers(get_chat_template("qwen3")) + + +# ── span-level masking tests ─────────────────────────────────────────── +# +# These render through transformers' generation tracker — the same code +# path `return_assistant_tokens_mask=True` (and therefore TRL's +# `assistant_only_loss=True`) uses — and assert on the exact substrings +# that end up in the loss. Character spans, not token spans, so no +# tokenizer download is needed. + + +def _assistant_spans(template_name: str, messages: list[dict]) -> list[str]: + """Return the substrings of the rendered conversation that are in the loss.""" + # transformers-internal: if this import moves, the public equivalent is + # `tokenizer.apply_chat_template(..., return_assistant_tokens_mask=True)`, + # which additionally requires a tokenizer. + chat_template_utils = pytest.importorskip("transformers.utils.chat_template_utils") + + compiled = chat_template_utils._compile_jinja_template(get_chat_template(template_name)) + rendered, indices = chat_template_utils._render_with_assistant_indices( + compiled, messages, None, None, False + ) + return [rendered[start:end] for start, end in indices] + + +def _think(reasoning: str, answer: str) -> str: + return f"\n{reasoning}\n\n\n{answer}" + + +def test_qwen3_excludes_history_assistant_turns() -> None: + """Qwen3 strips `` from assistant turns at or before the last + user query, so those turns render as a bare answer with no reasoning + trace. Training on them teaches the model to answer directly at + exactly the position where it always opens a `` block at + inference — a train/inference mismatch that trains against the + reasoning behaviour. Only the final exchange belongs in the loss. + """ + messages = [ + {"role": "user", "content": "Q1"}, + {"role": "assistant", "content": _think("R1", "A1")}, + {"role": "user", "content": "Q2"}, + {"role": "assistant", "content": _think("R2", "A2")}, + {"role": "user", "content": "Q3"}, + {"role": "assistant", "content": _think("R3", "A3")}, + ] + + assert _assistant_spans("qwen3", messages) == [_think("R3", "A3") + "<|im_end|>"] + + +def test_qwen3_keeps_every_assistant_turn_after_last_user_query() -> None: + """`ns.last_query_index` tracks the last *real* user message, not the + last message. A multi-step tool exchange therefore has several + assistant turns after it, and all of them keep their reasoning trace + when rendered — so all of them belong in the loss. + """ + messages = [ + {"role": "user", "content": "weather?"}, + { + "role": "assistant", + "content": "\nCall the tool.\n\n\n", + "tool_calls": [ + { + "type": "function", + "function": {"name": "get_weather", "arguments": {"city": "Berlin"}}, + } + ], + }, + {"role": "tool", "content": '{"temp_c": 18}'}, + {"role": "assistant", "content": "18C in Berlin."}, + ] + + spans = _assistant_spans("qwen3", messages) + assert len(spans) == 2 + assert spans[0].startswith("\nCall the tool.\n") + assert "" in spans[0] + assert spans[1].endswith("18C in Berlin.<|im_end|>") + + +def test_qwen3_excludes_prompt_header_and_turn_separator() -> None: + """The `<|im_start|>assistant\\n` header is prompt, and the `\\n` after + `<|im_end|>` is a turn separator — neither is something the model + generates, so both stay outside the span. + """ + messages = [ + {"role": "user", "content": "Q"}, + {"role": "assistant", "content": _think("R", "A")}, + ] + + (span,) = _assistant_spans("qwen3", messages) + assert not span.startswith("<|im_start|>") + assert span.endswith("<|im_end|>") + + +def test_olmo3_instruct_sft_masks_all_assistant_turns() -> None: + """Counterpoint to the qwen3 behaviour above, pinned deliberately: the + OLMo templates do not strip anything from history, so every assistant + turn renders exactly as generated and every one belongs in the loss. + The qwen3 exclusion is a Qwen-specific consequence of think-stripping, + not a house rule about multi-turn data. + """ + messages = [ + {"role": "user", "content": "Q1"}, + {"role": "assistant", "content": "A1"}, + {"role": "user", "content": "Q2"}, + {"role": "assistant", "content": "A2"}, + ] + + spans = _assistant_spans("olmo3-instruct-sft", messages) + assert len(spans) == 2 + assert spans[0].startswith("A1") + assert spans[1].startswith("A2") From 218858e03b44955824cb469dc90e91745d4b1b86 Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Fri, 7 Aug 2026 14:10:06 +0200 Subject: [PATCH 08/12] test: assert qwen3 template renders byte-identically to upstream Vendors the pristine Qwen/Qwen3-8B chat_template as a fixture and diffs rendered output across the 20-combination matrix. --- tests/fixtures/qwen3-8b-upstream.jinja | 89 ++++++++++++++ tests/test_chat_template_parity.py | 164 +++++++++++++++++++++++++ 2 files changed, 253 insertions(+) create mode 100644 tests/fixtures/qwen3-8b-upstream.jinja create mode 100644 tests/test_chat_template_parity.py diff --git a/tests/fixtures/qwen3-8b-upstream.jinja b/tests/fixtures/qwen3-8b-upstream.jinja new file mode 100644 index 0000000..01be9b3 --- /dev/null +++ b/tests/fixtures/qwen3-8b-upstream.jinja @@ -0,0 +1,89 @@ +{%- 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 XML tags:\n" }} + {%- for tool in tools %} + {{- "\n" }} + {{- tool | tojson }} + {%- endfor %} + {{- "\n\n\nFor each function call, return a json object with function name and arguments within XML tags:\n\n{\"name\": , \"arguments\": }\n<|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('') and message.content.endswith('')) %} + {%- 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 '' in content %} + {%- set reasoning_content = content.split('')[0].rstrip('\n').split('')[-1].lstrip('\n') %} + {%- set content = content.split('')[-1].lstrip('\n') %} + {%- endif %} + {%- endif %} + {%- if loop.index0 > ns.last_query_index %} + {%- if loop.last or (not loop.last and reasoning_content) %} + {{- '<|im_start|>' + message.role + '\n\n' + reasoning_content.strip('\n') + '\n\n\n' + content.lstrip('\n') }} + {%- else %} + {{- '<|im_start|>' + message.role + '\n' + content }} + {%- endif %} + {%- else %} + {{- '<|im_start|>' + message.role + '\n' + 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 %} + {{- '\n{"name": "' }} + {{- tool_call.name }} + {{- '", "arguments": ' }} + {%- if tool_call.arguments is string %} + {{- tool_call.arguments }} + {%- else %} + {{- tool_call.arguments | tojson }} + {%- endif %} + {{- '}\n' }} + {%- endfor %} + {%- endif %} + {{- '<|im_end|>\n' }} + {%- elif message.role == "tool" %} + {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %} + {{- '<|im_start|>user' }} + {%- endif %} + {{- '\n\n' }} + {{- content }} + {{- '\n' }} + {%- 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 %} + {{- '\n\n\n\n' }} + {%- endif %} +{%- endif %} \ No newline at end of file diff --git a/tests/test_chat_template_parity.py b/tests/test_chat_template_parity.py new file mode 100644 index 0000000..1b2f72e --- /dev/null +++ b/tests/test_chat_template_parity.py @@ -0,0 +1,164 @@ +"""Render-parity tests: registered templates vs. their pristine upstream. + +Splicing `{% generation %}` markers into a vendored chat template is only +safe if it changes *nothing* about what the template renders — the markers +carve up the loss mask, they must not move a single byte of the prompt. A +template that renders even slightly differently from the one the base model +was trained with is a silent train/inference skew. + +`tests/fixtures/qwen3-8b-upstream.jinja` is the unmodified `chat_template` +field from `Qwen/Qwen3-8B`'s `tokenizer_config.json` (4168 bytes, no +trailing newline; the repo publishes no standalone `chat_template.jinja`). +It is a reference copy — never edit it. To refresh it after an upstream +change, re-extract it verbatim and update `_UPSTREAM_SHA256`: + + python -c "import json,urllib.request as u; \\ + print(json.load(u.urlopen('https://huggingface.co/Qwen/Qwen3-8B/\\ +resolve/main/tokenizer_config.json'))['chat_template'], end='')" \\ + > tests/fixtures/qwen3-8b-upstream.jinja +""" + +from __future__ import annotations + +import hashlib +from pathlib import Path + +import pytest + +from post_training.chat_templates.registry import ( + get_chat_template, + has_generation_markers, +) + +_FIXTURES = Path(__file__).parent / "fixtures" +_UPSTREAM = _FIXTURES / "qwen3-8b-upstream.jinja" + +# Guards the reference copy against silent edits — including an editor or +# hook appending a trailing newline, which would change what it renders. +_UPSTREAM_SHA256 = "a55ee1b1660128b7098723e0abcd92caa0788061051c62d51cbe87d9cf1974d8" + + +# ── conversation matrix ──────────────────────────────────────────────── + +_TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string", "description": "City name."}}, + "required": ["city"], + }, + }, + } +] + +_CONVERSATIONS: dict[str, list[dict]] = { + "single-turn-reasoning": [ + {"role": "user", "content": "What is 17 * 23?"}, + {"role": "assistant", "content": "\n340 + 51.\n\n\n391."}, + ], + "multi-turn": [ + {"role": "user", "content": "Q1"}, + {"role": "assistant", "content": "\nR1\n\n\nA1"}, + {"role": "user", "content": "Q2"}, + {"role": "assistant", "content": "\nR2\n\n\nA2"}, + {"role": "user", "content": "Q3"}, + {"role": "assistant", "content": "\nR3\n\n\nA3"}, + ], + "no-think": [ + {"role": "user", "content": "Capital of France?"}, + {"role": "assistant", "content": "Paris."}, + {"role": "user", "content": "And Japan?"}, + {"role": "assistant", "content": "Tokyo."}, + ], + "tool-call": [ + {"role": "user", "content": "Weather in Berlin?"}, + { + "role": "assistant", + "content": "\nCall the tool.\n\n\n", + "tool_calls": [ + { + "type": "function", + "function": {"name": "get_weather", "arguments": {"city": "Berlin"}}, + } + ], + }, + {"role": "tool", "content": '{"temp_c": 18, "sky": "clear"}'}, + {"role": "assistant", "content": "18C and clear."}, + ], + "system+tools": [ + {"role": "system", "content": "You are a terse assistant."}, + {"role": "user", "content": "Weather in Berlin?"}, + {"role": "assistant", "content": "\nLooks clear.\n\n\nSunny."}, + ], +} + + +def _render(template: str, messages: list[dict], tools: list[dict] | None, agp: bool) -> str: + """Render *template* through transformers' own chat-template machinery. + + Going through transformers rather than a bare `jinja2.Environment` + matters: the parity claim is about what `apply_chat_template` actually + produces, which depends on `trim_blocks`/`lstrip_blocks` and on the + `{% generation %}` extension being registered. + """ + # transformers-internal. The public equivalent is + # `tokenizer.apply_chat_template(..., tokenize=False)`, which would drag + # a tokenizer download into what is otherwise a pure-Jinja test. + utils = pytest.importorskip("transformers.utils.chat_template_utils") + + rendered, _ = utils._render_with_assistant_indices( + utils._compile_jinja_template(template), messages, tools, None, agp + ) + return rendered + + +# ── fixture integrity ────────────────────────────────────────────────── + + +def test_upstream_fixture_is_unmodified() -> None: + """The fixture is only meaningful as an *untouched* reference copy.""" + digest = hashlib.sha256(_UPSTREAM.read_bytes()).hexdigest() + assert digest == _UPSTREAM_SHA256, ( + "tests/fixtures/qwen3-8b-upstream.jinja has been modified. It must stay " + "byte-identical to Qwen/Qwen3-8B's upstream chat_template; see this " + "module's docstring for how to refresh it." + ) + + +def test_upstream_fixture_has_no_generation_markers() -> None: + """Confirms the fixture really is the pristine upstream template: the + markers are ours, so their absence here is what makes the render-parity + comparison below meaningful. + """ + assert not has_generation_markers(_UPSTREAM.read_text()) + assert has_generation_markers(get_chat_template("qwen3")) + + +# ── render parity ────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("shape", sorted(_CONVERSATIONS)) +@pytest.mark.parametrize("add_generation_prompt", [False, True]) +@pytest.mark.parametrize("with_tools", [False, True]) +def test_qwen3_renders_identically_to_upstream( + shape: str, add_generation_prompt: bool, with_tools: bool +) -> None: + """Our marker-spliced qwen3 template must render byte-identically to the + pristine upstream template in every combination — assistant content, + reasoning traces, tool calls, the tools system block, and the trailing + generation prompt all included. + + If this fails, the marker splice moved a byte of the prompt and the + template no longer matches what Qwen3 was trained with. + """ + messages = _CONVERSATIONS[shape] + tools = _TOOLS if with_tools else None + + ours = _render(get_chat_template("qwen3"), messages, tools, add_generation_prompt) + upstream = _render(_UPSTREAM.read_text(), messages, tools, add_generation_prompt) + + assert ours == upstream From e31a86064d1869ba0606b531325cd2709421b1f5 Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Fri, 7 Aug 2026 14:24:01 +0200 Subject: [PATCH 09/12] refactor: rename qwen3 jinja file --- src/post_training/chat_templates/registry.py | 2 +- .../chat_templates/templates/{qwen3-8b.jinja => qwen3.jinja} | 0 2 files changed, 1 insertion(+), 1 deletion(-) rename src/post_training/chat_templates/templates/{qwen3-8b.jinja => qwen3.jinja} (100%) diff --git a/src/post_training/chat_templates/registry.py b/src/post_training/chat_templates/registry.py index 2a8ddc5..c067e2d 100644 --- a/src/post_training/chat_templates/registry.py +++ b/src/post_training/chat_templates/registry.py @@ -62,7 +62,7 @@ "olmo3-think-sft": "olmo3-think-sft.jinja", "apertus": "apertus.jinja", "tulu3": "tulu3.jinja", - "qwen3": "qwen3-8b.jinja", + "qwen3": "qwen3.jinja", } diff --git a/src/post_training/chat_templates/templates/qwen3-8b.jinja b/src/post_training/chat_templates/templates/qwen3.jinja similarity index 100% rename from src/post_training/chat_templates/templates/qwen3-8b.jinja rename to src/post_training/chat_templates/templates/qwen3.jinja From db4f5e8fb6342457849559a5707f6712f8ba7fe7 Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Wed, 19 Aug 2026 17:45:07 +0200 Subject: [PATCH 10/12] feat: filter out rows with all zero assistant loss mask --- scripts/data.py | 13 ++- src/post_training/data/loader.py | 22 ++--- src/post_training/methods/dpo.py | 13 ++- src/post_training/methods/sft.py | 135 +++++++++++++++++++++++++++++-- 4 files changed, 161 insertions(+), 22 deletions(-) diff --git a/scripts/data.py b/scripts/data.py index 1f0f49f..94b6ee7 100755 --- a/scripts/data.py +++ b/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, ) diff --git a/src/post_training/data/loader.py b/src/post_training/data/loader.py index 3f81258..4d38677 100644 --- a/src/post_training/data/loader.py +++ b/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) diff --git a/src/post_training/methods/dpo.py b/src/post_training/methods/dpo.py index 42f74c3..4b2a0ce 100644 --- a/src/post_training/methods/dpo.py +++ b/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), diff --git a/src/post_training/methods/sft.py b/src/post_training/methods/sft.py index 8e859ce..2f3bc15 100644 --- a/src/post_training/methods/sft.py +++ b/src/post_training/methods/sft.py @@ -3,11 +3,13 @@ from __future__ import annotations import logging +from functools import partial from pathlib import Path from typing import TYPE_CHECKING from accelerate import PartialState -from datasets import Features, List, Value +from datasets import Dataset, Features, List, Value +from transformers import PreTrainedTokenizerBase from trl import SFTConfig, SFTTrainer from post_training.chat_templates.registry import has_generation_markers @@ -26,6 +28,9 @@ logger = logging.getLogger(__name__) +# Warn instead of info once this share of rows is dropped. +_DROP_WARN_PERCENT = 1.0 + MESSAGES_FEATURES = Features( { "messages": List( @@ -38,9 +43,117 @@ ) -def _sft_row_filter(example: dict) -> bool: - """Keep only rows with a non-empty ``messages`` list.""" - return len(example.get("messages", [])) > 0 +def _filter_sft_rows( + ds: Dataset, + num_proc: int, + *, + tokenizer: PreTrainedTokenizerBase, + max_length: int | None = None, +) -> Dataset: + """Drop rows that contribute nothing to the SFT loss. + + A row is dropped when its ``messages`` list is empty, or when rendering it + through the chat template gives an all-zero ``assistant_masks``. The second + case happens when a conversation ends on a non-assistant turn, or when the + template strips the only assistant content the row holds. Such rows cost + compute and produce a zero gradient. + + Parameters + ---------- + ds: + Dataset with a ``messages`` column. + num_proc: + Number of worker processes for ``.filter()`` and ``.map()``. + tokenizer: + Tokenizer whose ``chat_template`` carries ``{% generation %}`` markers. + Bind it at the call site with :func:`functools.partial`. + max_length: + Optional truncation length that must match the trainer's own truncation, + so a row whose assistant tokens fall outside the window is dropped too. + Leave it ``None`` only if the trainer keeps every token of every row. + + Returns + ------- + datasets.Dataset + The rows that stay. + + Raises + ------ + ValueError + If no row survives either stage. + """ + ds = ds.filter( + lambda row: len(row["messages"]) > 0, num_proc=num_proc, desc="filtering empty messages" + ) + + if len(ds) == 0: + raise ValueError( + "No rows remain after dropping empty 'messages' lists. " + "Check the dataset and its transform." + ) + + template_kwargs = {"truncation": True, "max_length": max_length} if max_length else {} + + # remove_columns keeps the map cache tiny: it holds the flag alone, not a copy of the data. + in_loss = ds.map( + lambda row: { + "in_loss": any( + tokenizer.apply_chat_template( + row["messages"], + return_dict=True, + return_assistant_tokens_mask=True, + **template_kwargs, + )["assistant_masks"] + ) + }, + num_proc=num_proc, + remove_columns=ds.column_names, + desc="computing assistant loss masks", + )["in_loss"] + + keep: list[int] = [] + dropped: list[int] = [] + for index, ok in enumerate(in_loss): + if ok: + keep.append(index) + else: + dropped.append(index) + + dropped_percent = 100 * len(dropped) / len(ds) + summary = ( + f"{len(ds)} rows, {len(dropped)} with an all-zero assistant mask ({dropped_percent:.3f}%)." + ) + + if not keep: + cause = ( + "The chat template is probably missing the {% generation %}…{% endgeneration %} " + "markers around the assistant content, which makes every mask all-zero." + ) + if max_length is not None: + cause += ( + f" A max_length of {max_length} may also be short enough to truncate " + "away every assistant turn." + ) + raise ValueError(f"{summary} Every row would contribute zero loss. {cause}") + + if dropped_percent > _DROP_WARN_PERCENT: + logger.warning(summary) + else: + logger.info(summary) + + if dropped: + patterns = { + " -> ".join(message["role"] for message in row["messages"][-3:]) + for row in ds.select(dropped) + } + logger.info( + "Role sequences of the %d dropped rows (%d unique):\n%s", + len(dropped), + len(patterns), + "\n".join(sorted(patterns)), + ) + + return ds.select(keep) def build_sft_trainer(config: PostTrainingConfig, run_dir: Path) -> SFTTrainer: @@ -92,7 +205,19 @@ def build_sft_trainer(config: PostTrainingConfig, run_dir: Path) -> SFTTrainer: with PartialState().main_process_first(): dataset = load_and_mix_datasets( config.data, - row_filter=_sft_row_filter, + dataset_filter_fn=partial( + _filter_sft_rows, + tokenizer=tokenizer, + # Both TRL paths cut each row at max_length, so the filter must + # measure the same window. Without packing, TRL truncates the + # row. With packing, the default "bfd" strategy keeps only the + # first max_length fragment and discards the overflow + # (trl/data_utils.py _pack_bfd). Revisit if packing_strategy + # ever becomes configurable: "bfd-requeue" and "wrapped" both + # keep the overflow, so truncation-based filtering would then + # drop rows the trainer still trains on. + max_length=mc.max_seq_length, + ), columns_to_keep=["messages"], features=MESSAGES_FEATURES, ) From 189bb7bf9d449fe348abd7d4875e9d9e6ee7a826 Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Wed, 19 Aug 2026 17:46:10 +0200 Subject: [PATCH 11/12] tests: add tests for row filters --- tests/test_row_filters.py | 385 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 385 insertions(+) create mode 100644 tests/test_row_filters.py diff --git a/tests/test_row_filters.py b/tests/test_row_filters.py new file mode 100644 index 0000000..666ba88 --- /dev/null +++ b/tests/test_row_filters.py @@ -0,0 +1,385 @@ +"""Tests for the method-specific dataset filters. + +`_filter_sft_rows` drops rows that would contribute nothing to the SFT loss: +an empty `messages` list, or a conversation whose rendered `assistant_masks` +is all zero. A row that survives to the trainer but produces no gradient +costs compute and quietly shrinks the effective dataset, so the filter is +worth pinning down. + +`_filter_dpo_rows` drops preference pairs missing a `chosen` or a `rejected` +side. It reads those two columns; a DPO dataset has no `messages` column at +all, so a filter that reaches for one raises `KeyError` on every run. + +Which rows an all-zero mask catches depends on the template, so these tests +use two marker-bearing templates deliberately: + +* `qwen3` puts only the assistant turns **after the last user query** in the + loss. It strips `` from earlier turns, so those render as a bare + answer and are excluded. A conversation that ends on a user turn therefore + has an all-zero mask under qwen3 and gets dropped. +* `olmo3-instruct-sft` puts **every** assistant turn in the loss, so the same + conversation survives. + +`tests/test_chat_template_masking.py` pins that span-level behaviour itself. +Here it is the input to the filter. + +No tokenizer is downloaded. `_CharTokenizer` renders the real registered +templates through transformers' own generation tracker — the same code path +`return_assistant_tokens_mask=True` uses — and then treats one character as +one token. Character offsets map one-to-one onto token offsets, so the masks +and the `max_length` cuts carry the same meaning as a real tokenizer's. +""" + +from __future__ import annotations + +import logging + +import pytest +from datasets import Dataset + +from post_training.chat_templates.registry import get_chat_template +from post_training.methods.dpo import _filter_dpo_rows +from post_training.methods.sft import _filter_sft_rows + +SFT_LOGGER = "post_training.methods.sft" + +# Loss on the assistant turns after the last user query only. +MARKED_TEMPLATE = "qwen3" +# Loss on every assistant turn — the contrast case. +HISTORY_TEMPLATE = "olmo3-instruct-sft" +# No `{% generation %}` markers at all — every mask comes back all zero. +UNMARKED_TEMPLATE = "olmo3" + +_COMPILED: dict[str, object] = {} + + +class _CharTokenizer: + """Stand-in for a HF tokenizer, with one character per token. + + Holds the template *name* rather than the compiled Jinja object so that + `datasets` can still fingerprint the `.map()` closure that captures it. + """ + + def __init__(self, template_name: str) -> None: + self.template_name = template_name + + @property + def chat_template(self) -> str: + return get_chat_template(self.template_name) + + def _render(self, conversation: list[dict]) -> tuple[str, list[tuple[int, int]]]: + utils = pytest.importorskip("transformers.utils.chat_template_utils") + if self.template_name not in _COMPILED: + _COMPILED[self.template_name] = utils._compile_jinja_template(self.chat_template) + return utils._render_with_assistant_indices( + _COMPILED[self.template_name], conversation, None, None, False + ) + + def render(self, conversation: list[dict]) -> str: + rendered, _ = self._render(conversation) + return rendered + + def apply_chat_template( + self, + conversation: list[dict], + *, + return_dict: bool = False, + return_assistant_tokens_mask: bool = False, + truncation: bool = False, + max_length: int | None = None, + ) -> dict: + rendered, indices = self._render(conversation) + + input_ids = [ord(character) for character in rendered] + assistant_masks = [0] * len(rendered) + for start, end in indices: + for position in range(start, end): + assistant_masks[position] = 1 + + if truncation and max_length is not None: + input_ids = input_ids[:max_length] + assistant_masks = assistant_masks[:max_length] + + return {"input_ids": input_ids, "assistant_masks": assistant_masks} + + +@pytest.fixture +def tokenizer() -> _CharTokenizer: + """A qwen3 tokenizer — the template our Qwen3 SFT config uses.""" + return _CharTokenizer(MARKED_TEMPLATE) + + +def _exchange(index: int = 0) -> list[dict]: + return [ + {"role": "user", "content": f"question-{index}"}, + {"role": "assistant", "content": f"answer-{index}"}, + ] + + +def _sft_dataset(rows: list[list[dict]]) -> Dataset: + return Dataset.from_dict({"messages": rows}) + + +# ── SFT: which rows survive ──────────────────────────────────────────── + + +def test_keeps_rows_with_an_assistant_turn(tokenizer) -> None: + ds = _sft_dataset([_exchange(0), _exchange(1)]) + + kept = _filter_sft_rows(ds, num_proc=1, tokenizer=tokenizer) + + assert len(kept) == 2 + + +def test_drops_rows_with_an_empty_messages_list(tokenizer) -> None: + ds = _sft_dataset([_exchange(0), [], _exchange(1)]) + + kept = _filter_sft_rows(ds, num_proc=1, tokenizer=tokenizer) + + assert kept["messages"] == [_exchange(0), _exchange(1)] + + +def test_drops_a_conversation_with_no_assistant_turn(tokenizer) -> None: + """A prompt-only row renders to an all-zero mask, so it trains on nothing.""" + ds = _sft_dataset([[{"role": "user", "content": "orphan question"}], _exchange(0)]) + + kept = _filter_sft_rows(ds, num_proc=1, tokenizer=tokenizer) + + assert kept["messages"] == [_exchange(0)] + + +def test_keeps_a_multi_turn_row_when_only_the_final_turn_is_in_the_loss(tokenizer) -> None: + """qwen3 excludes the assistant turns at or before the last user query, so + this row contributes one turn out of two. One is enough — the filter asks + whether *any* token enters the loss, not whether every turn does. + """ + multi_turn = [*_exchange(0), *_exchange(1)] + ds = _sft_dataset([multi_turn]) + + kept = _filter_sft_rows(ds, num_proc=1, tokenizer=tokenizer) + + assert kept["messages"] == [multi_turn] + + +def test_drops_a_conversation_ending_on_a_user_turn_under_qwen3(tokenizer) -> None: + """Under qwen3 every assistant turn here sits at or before the last user + query, so all of them render think-stripped and stay out of the loss. The + mask is all zero and the row trains on nothing. + + This is the case the filter exists for. Such a row looks perfectly valid + — it has an assistant turn and non-empty content — so nothing before the + mask computation catches it. + """ + trailing_user = [*_exchange(0), {"role": "user", "content": "follow-up"}] + ds = _sft_dataset([trailing_user, _exchange(1)]) + + kept = _filter_sft_rows(ds, num_proc=1, tokenizer=tokenizer) + + assert kept["messages"] == [_exchange(1)] + + +def test_keeps_a_conversation_ending_on_a_user_turn_under_olmo() -> None: + """Counterpoint to the qwen3 case above, pinned deliberately. The OLMo + templates keep every assistant turn in the loss, so the identical row + survives. Which rows the filter drops is a property of the template, not + a house rule about trailing user turns. + """ + trailing_user = [*_exchange(0), {"role": "user", "content": "follow-up"}] + ds = _sft_dataset([trailing_user]) + + kept = _filter_sft_rows(ds, num_proc=1, tokenizer=_CharTokenizer(HISTORY_TEMPLATE)) + + assert kept["messages"] == [trailing_user] + + +def test_preserves_row_order_and_content(tokenizer) -> None: + rows = [_exchange(0), [], _exchange(1), [], _exchange(2)] + ds = _sft_dataset(rows) + + kept = _filter_sft_rows(ds, num_proc=1, tokenizer=tokenizer) + + assert kept["messages"] == [_exchange(0), _exchange(1), _exchange(2)] + + +# ── SFT: truncation parity with the trainer ──────────────────────────── + + +def test_drops_a_row_whose_assistant_turn_truncation_removes(tokenizer) -> None: + """TRL cuts every row at `max_length` — plain truncation without packing, + and the first-fragment-only rule of the default `bfd` packing strategy. + A row whose assistant turn sits past that cut trains on nothing, so the + filter must measure the same window. + """ + long_prompt = [ + {"role": "user", "content": "x" * 500}, + {"role": "assistant", "content": "answer"}, + ] + ds = _sft_dataset([long_prompt]) + + assert len(_filter_sft_rows(ds, num_proc=1, tokenizer=tokenizer)) == 1 + + with pytest.raises(ValueError, match="zero loss"): + _filter_sft_rows(ds, num_proc=1, tokenizer=tokenizer, max_length=32) + + +def test_keeps_a_row_whose_assistant_turn_fits_inside_max_length(tokenizer) -> None: + ds = _sft_dataset([_exchange(0)]) + width = len(tokenizer.render(_exchange(0))) + + kept = _filter_sft_rows(ds, num_proc=1, tokenizer=tokenizer, max_length=width) + + assert len(kept) == 1 + + +# ── SFT: the failure modes that must be loud ─────────────────────────── + + +def test_raises_when_every_row_has_an_empty_messages_list(tokenizer) -> None: + ds = _sft_dataset([[], []]) + + with pytest.raises(ValueError, match="empty 'messages'"): + _filter_sft_rows(ds, num_proc=1, tokenizer=tokenizer) + + +def test_raises_when_the_template_has_no_generation_markers() -> None: + """Without `{% generation %}` markers every mask is all zero. The filter + would otherwise return an empty dataset and the failure would surface far + downstream as a confusing weight or resampling error. + """ + ds = _sft_dataset([_exchange(0), _exchange(1)]) + + with pytest.raises(ValueError, match="generation"): + _filter_sft_rows(ds, num_proc=1, tokenizer=_CharTokenizer(UNMARKED_TEMPLATE)) + + +def test_raises_when_every_row_ends_on_a_user_turn_under_qwen3(tokenizer) -> None: + """A whole mix of trailing-user conversations is a plausible mistake — a + prompt-completion dataset transformed into `messages` without pairing the + completion back in. Under qwen3 it yields zero usable rows, and the run + must stop with that stated rather than fail later on an empty mix. + """ + ds = _sft_dataset([[*_exchange(i), {"role": "user", "content": "next"}] for i in range(3)]) + + with pytest.raises(ValueError, match="zero loss"): + _filter_sft_rows(ds, num_proc=1, tokenizer=tokenizer) + + +def test_total_drop_error_names_max_length_only_when_it_is_set(tokenizer) -> None: + """Truncation and missing markers both empty the dataset. The message + must name the cause that applies, so the reader is not sent to the wrong + one. + """ + ds = _sft_dataset([[{"role": "user", "content": "no assistant turn"}]]) + + with pytest.raises(ValueError) as without_length: + _filter_sft_rows(ds, num_proc=1, tokenizer=tokenizer) + assert "max_length" not in str(without_length.value) + + with pytest.raises(ValueError) as with_length: + _filter_sft_rows(ds, num_proc=1, tokenizer=tokenizer, max_length=4096) + assert "max_length of 4096" in str(with_length.value) + + +def test_requires_the_tokenizer_as_a_keyword_argument() -> None: + """`build_sft_trainer` binds the tokenizer with `functools.partial`. A + missing tokenizer must fail at the call, not as an `AttributeError` inside + the worker pool — and not silently under `python -O`, which strips asserts. + """ + ds = _sft_dataset([_exchange(0)]) + + with pytest.raises(TypeError, match="tokenizer"): + _filter_sft_rows(ds, num_proc=1) + + +# ── SFT: reporting ───────────────────────────────────────────────────── + + +def test_warns_when_the_drop_rate_passes_the_threshold(tokenizer, caplog) -> None: + prompt_only = [{"role": "user", "content": "orphan"}] + ds = _sft_dataset([_exchange(i) for i in range(9)] + [prompt_only]) + + with caplog.at_level(logging.INFO, logger=SFT_LOGGER): + _filter_sft_rows(ds, num_proc=1, tokenizer=tokenizer) + + levels = {record.levelno for record in caplog.records if "all-zero" in record.message} + assert levels == {logging.WARNING} + + +def test_logs_at_info_when_the_drop_rate_stays_under_the_threshold(tokenizer, caplog) -> None: + prompt_only = [{"role": "user", "content": "orphan"}] + ds = _sft_dataset([_exchange(i) for i in range(199)] + [prompt_only]) + + with caplog.at_level(logging.INFO, logger=SFT_LOGGER): + _filter_sft_rows(ds, num_proc=1, tokenizer=tokenizer) + + levels = {record.levelno for record in caplog.records if "all-zero" in record.message} + assert levels == {logging.INFO} + + +def test_reports_each_distinct_role_sequence_once(tokenizer, caplog) -> None: + """The dropped-row report exists to show *why* rows went away. It lists + the trailing role sequence of every dropped row, de-duplicated, so a + million identical drops still print one line. + + Under qwen3 the `user -> assistant -> user` line is the one that matters: + it names the think-stripping exclusion, which is the non-obvious reason a + complete-looking conversation contributes nothing. + """ + prompt_only = [{"role": "user", "content": "orphan"}] + trailing_user = [*_exchange(0), {"role": "user", "content": "follow-up"}] + ds = _sft_dataset([_exchange(1), prompt_only, prompt_only, prompt_only, trailing_user]) + + with caplog.at_level(logging.INFO, logger=SFT_LOGGER): + _filter_sft_rows(ds, num_proc=1, tokenizer=tokenizer) + + report = next(record for record in caplog.records if "Role sequences" in record.getMessage()) + body = report.getMessage().split("\n", 1)[1] + + assert "4 dropped rows (2 unique)" in report.getMessage() + assert sorted(body.splitlines()) == ["user", "user -> assistant -> user"] + + +# ── DPO ──────────────────────────────────────────────────────────────── + + +def _dpo_dataset(pairs: list[tuple[list, list]]) -> Dataset: + return Dataset.from_dict( + { + "chosen": [chosen for chosen, _ in pairs], + "rejected": [rejected for _, rejected in pairs], + } + ) + + +def test_dpo_keeps_complete_preference_pairs() -> None: + ds = _dpo_dataset([(_exchange(0), _exchange(1)), (_exchange(2), _exchange(3))]) + + assert len(_filter_dpo_rows(ds, num_proc=1)) == 2 + + +def test_dpo_drops_a_pair_with_an_empty_chosen_side() -> None: + ds = _dpo_dataset([([], _exchange(0)), (_exchange(1), _exchange(2))]) + + kept = _filter_dpo_rows(ds, num_proc=1) + + assert kept["chosen"] == [_exchange(1)] + + +def test_dpo_drops_a_pair_with_an_empty_rejected_side() -> None: + ds = _dpo_dataset([(_exchange(0), []), (_exchange(1), _exchange(2))]) + + kept = _filter_dpo_rows(ds, num_proc=1) + + assert kept["rejected"] == [_exchange(2)] + + +def test_dpo_reads_chosen_and_rejected_rather_than_messages() -> None: + """Regression guard. A DPO dataset carries `chosen` / `rejected` and no + `messages` column (see `configs/trl/dpo.yaml`, which uses `transform: null` + on an already-paired dataset). A filter that reads `messages` raises + `KeyError` on every DPO run. + """ + ds = _dpo_dataset([(_exchange(0), _exchange(1))]) + assert "messages" not in ds.column_names + + assert len(_filter_dpo_rows(ds, num_proc=1)) == 1 From 604e7312f4a1b9a61b0f0996480333f0abb758c0 Mon Sep 17 00:00:00 2001 From: Konstantin Nikolaou <87869540+KonstiNik@users.noreply.github.com> Date: Tue, 25 Aug 2026 12:00:34 +0200 Subject: [PATCH 12/12] fix: warn when max_length cuts a supervised span part-way (#49) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: warn when max_length cuts a supervised span part-way `any(assistant_masks)` only catches rows whose assistant turn is truncated away entirely. A row that straddles max_length keeps a non-zero mask, so it stays in training and teaches an answer cut mid-sentence with no end-of-turn token. One untruncated render classifies all three cases instead — no assistant tokens, all of them past the cap, or straddling it — at the same cost. The keep/drop decision is unchanged. * feat: add sft.truncated_span_action to drop rows a cut span damages Warning is right when the cap is wrong, since raising max_seq_length keeps the data. It is not enough when the sequence length is fixed by memory and dropping is the only lever left, so make it a switch. Default stays warn. Drop still warns. data.datasets[].weight multiplies the rows that survive filtering, so an uneven drop changes a dataset's share of the mixture while the config still claims the original weights — that has to be in the log rather than inferred from a row count. --------- Co-authored-by: Konstantin Nikolaou --- configs/trl/sft.yaml | 5 + src/post_training/config.py | 7 ++ src/post_training/methods/sft.py | 191 +++++++++++++++++++++++++------ tests/test_row_filters.py | 177 ++++++++++++++++++++++++++++ 4 files changed, 345 insertions(+), 35 deletions(-) diff --git a/configs/trl/sft.yaml b/configs/trl/sft.yaml index 5699cf2..415e58e 100644 --- a/configs/trl/sft.yaml +++ b/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: diff --git a/src/post_training/config.py b/src/post_training/config.py index fb65058..01734bf 100644 --- a/src/post_training/config.py +++ b/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 diff --git a/src/post_training/methods/sft.py b/src/post_training/methods/sft.py index 2f3bc15..9fe5e59 100644 --- a/src/post_training/methods/sft.py +++ b/src/post_training/methods/sft.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +from collections import Counter from functools import partial from pathlib import Path from typing import TYPE_CHECKING @@ -31,6 +32,63 @@ # Warn instead of info once this share of rows is dropped. _DROP_WARN_PERCENT = 1.0 +# What `max_length` does to one row's supervised span. +_KEEP = 0 # the whole supervised span survives +_CUT = 1 # kept, but max_length cuts the span part-way through +_NO_ASSISTANT = 2 # dropped: no assistant tokens anywhere in the row +_BEYOND_CAP = 3 # dropped: assistant tokens exist, all of them past max_length + +#: valid values for SFTMethodConfig.truncated_span_action +_TRUNCATED_SPAN_ACTIONS = ("warn", "drop") + + +def _classify_row( + messages: list[dict], + tokenizer: PreTrainedTokenizerBase, + max_length: int | None, +) -> int: + """Say what ``max_length`` does to one row's supervised span. + + One untruncated render answers all three questions, so this costs the same + as the plain ``any(assistant_masks)`` check and distinguishes three cases a + single boolean conflates: + + * no assistant tokens at all — the row's SHAPE is wrong, or the template + excludes every turn it has + * assistant tokens exist but lie entirely past ``max_length`` — the CAP is + wrong, and the row would be fine at a larger one + * the span starts inside the window and continues past it — the row stays in + training and teaches an answer that stops mid-sentence + + The third is the dangerous one: under final-turn-only supervision the + assistant content sits at the END of the row, so a row only slightly over + ``max_length`` keeps a non-zero mask, passes any ``any(...)`` check, and + trains on a truncated trace with no end-of-turn token. Nothing downstream + reports it. + + The render must be UNTRUNCATED: a truncated mask stops at ``max_length``, so + the third case would be invisible. The window is then applied by slicing, + which is right-truncation by construction — what both TRL paths do (the + non-packing path truncates the row; ``bfd`` keeps the first fragment). + """ + rendered = tokenizer.apply_chat_template( + messages, + return_dict=True, + return_assistant_tokens_mask=True, + ) + mask = rendered["assistant_masks"] + + if not any(mask): + return _NO_ASSISTANT + if max_length is None: + return _KEEP + if not any(mask[:max_length]): + return _BEYOND_CAP + if any(mask[max_length:]): + return _CUT + return _KEEP + + MESSAGES_FEATURES = Features( { "messages": List( @@ -49,14 +107,22 @@ def _filter_sft_rows( *, tokenizer: PreTrainedTokenizerBase, max_length: int | None = None, + truncated_span_action: str = "warn", ) -> Dataset: - """Drop rows that contribute nothing to the SFT loss. + """Drop rows that contribute nothing to the SFT loss, and report what does. A row is dropped when its ``messages`` list is empty, or when rendering it - through the chat template gives an all-zero ``assistant_masks``. The second - case happens when a conversation ends on a non-assistant turn, or when the - template strips the only assistant content the row holds. Such rows cost - compute and produce a zero gradient. + through the chat template leaves no assistant token inside the ``max_length`` + window. The second case happens when a conversation ends on a non-assistant + turn, when the template strips the only assistant content the row holds, or + when the assistant turn sits entirely past the cap. Such rows cost compute + and produce a zero gradient. + + Dropping is only half the job. A row whose supervised span STARTS inside the + window and continues past it is kept — it does produce gradient — but it + teaches an answer that stops mid-sentence with no end-of-turn token. That is + invisible to a keep/drop count, so it is reported separately and loudly; see + :func:`_classify_row`. Parameters ---------- @@ -71,6 +137,13 @@ def _filter_sft_rows( Optional truncation length that must match the trainer's own truncation, so a row whose assistant tokens fall outside the window is dropped too. Leave it ``None`` only if the trainer keeps every token of every row. + truncated_span_action: + ``"warn"`` (default) keeps a row whose supervised span straddles + ``max_length`` and reports it; ``"drop"`` removes it, for when the + sequence length is fixed by memory and shortening the data is the only + lever left. Drop still warns: it is not the quiet option, because + ``data.datasets[].weight`` multiplies the rows that survive this filter, + so an uneven drop changes the mixture the config asked for. Returns ------- @@ -80,8 +153,15 @@ def _filter_sft_rows( Raises ------ ValueError - If no row survives either stage. + If no row survives either stage, or ``truncated_span_action`` is not a + recognised value. """ + if truncated_span_action not in _TRUNCATED_SPAN_ACTIONS: + raise ValueError( + f"sft.truncated_span_action must be one of " + f"{', '.join(_TRUNCATED_SPAN_ACTIONS)}; got {truncated_span_action!r}." + ) + ds = ds.filter( lambda row: len(row["messages"]) > 0, num_proc=num_proc, desc="filtering empty messages" ) @@ -92,48 +172,52 @@ def _filter_sft_rows( "Check the dataset and its transform." ) - template_kwargs = {"truncation": True, "max_length": max_length} if max_length else {} - - # remove_columns keeps the map cache tiny: it holds the flag alone, not a copy of the data. - in_loss = ds.map( - lambda row: { - "in_loss": any( - tokenizer.apply_chat_template( - row["messages"], - return_dict=True, - return_assistant_tokens_mask=True, - **template_kwargs, - )["assistant_masks"] - ) - }, + # remove_columns keeps the map cache tiny: it holds the verdict alone, not a copy of the data. + verdicts = ds.map( + lambda row: {"verdict": _classify_row(row["messages"], tokenizer, max_length)}, num_proc=num_proc, remove_columns=ds.column_names, desc="computing assistant loss masks", - )["in_loss"] + )["verdict"] + counts = Counter(verdicts) + cut_is_kept = truncated_span_action == "warn" + kept_verdicts = (_KEEP, _CUT) if cut_is_kept else (_KEEP,) keep: list[int] = [] dropped: list[int] = [] - for index, ok in enumerate(in_loss): - if ok: - keep.append(index) - else: - dropped.append(index) + for index, verdict in enumerate(verdicts): + (keep if verdict in kept_verdicts else dropped).append(index) dropped_percent = 100 * len(dropped) / len(ds) + # "within the length window" only when there IS one — the total-drop error + # below must not name max_length when it was never set. + scope = " within the length window" if max_length is not None else "" summary = ( - f"{len(ds)} rows, {len(dropped)} with an all-zero assistant mask ({dropped_percent:.3f}%)." + f"{len(ds)} rows, {len(dropped)} with an all-zero assistant mask{scope} " + f"({dropped_percent:.3f}%)." ) if not keep: - cause = ( - "The chat template is probably missing the {% generation %}…{% endgeneration %} " - "markers around the assistant content, which makes every mask all-zero." - ) - if max_length is not None: - cause += ( - f" A max_length of {max_length} may also be short enough to truncate " - "away every assistant turn." + if counts[_BEYOND_CAP] and not counts[_NO_ASSISTANT]: + # Every row HAS an assistant turn, so the template is fine and the + # cap is the whole story. Sending the reader to the template here + # would cost them a long detour. + cause = ( + f"Every row has assistant tokens, but in every row all of them fall " + f"beyond max_length={max_length}. This is a max_seq_length problem, " + f"not a template problem: raise sft.max_seq_length, or shorten the " + f"rows upstream." ) + else: + cause = ( + "The chat template is probably missing the {% generation %}…{% endgeneration %} " + "markers around the assistant content, which makes every mask all-zero." + ) + if max_length is not None: + cause += ( + f" A max_length of {max_length} may also be short enough to truncate " + "away every assistant turn." + ) raise ValueError(f"{summary} Every row would contribute zero loss. {cause}") if dropped_percent > _DROP_WARN_PERCENT: @@ -141,6 +225,42 @@ def _filter_sft_rows( else: logger.info(summary) + # Warned about in BOTH modes, at any rate above zero. Under "warn" the rows are + # kept, so no drop count can reveal them. Under "drop" they are gone, which is + # not quiet either: data.datasets[].weight multiplies surviving rows, so an + # uneven drop silently changes this dataset's share of the mixture. + if counts[_CUT]: + logger.warning( + "%d of %d rows (%.3f%%) are cut by max_length=%d PART-WAY THROUGH their " + "supervised span — each teaches an answer that stops mid-sentence with no " + "end-of-turn token, an unterminated block for reasoning data. %s " + "Raising sft.max_seq_length above the rows' length is the fix that keeps " + "the data.", + counts[_CUT], + len(ds), + 100 * counts[_CUT] / len(ds), + max_length, + ( + "They STAY in training (sft.truncated_span_action='warn')." + if cut_is_kept + else "They were DROPPED (sft.truncated_span_action='drop'), which " + "shifts this dataset's share of the mixture: data.datasets[].weight " + "applies to the rows that survive filtering." + ), + ) + + if max_length is not None and (counts[_NO_ASSISTANT] or counts[_BEYOND_CAP]): + # Which of the two reasons a row was dropped decides what to change: + # the data and template for the first, sft.max_seq_length for the second. + logger.info( + "Dropped rows by cause: %d with no assistant tokens at all, " + "%d with every assistant token beyond max_length=%d%s.", + counts[_NO_ASSISTANT], + counts[_BEYOND_CAP], + max_length, + "" if cut_is_kept else f", {counts[_CUT]} with a span cut part-way", + ) + if dropped: patterns = { " -> ".join(message["role"] for message in row["messages"][-3:]) @@ -217,6 +337,7 @@ def build_sft_trainer(config: PostTrainingConfig, run_dir: Path) -> SFTTrainer: # keep the overflow, so truncation-based filtering would then # drop rows the trainer still trains on. max_length=mc.max_seq_length, + truncated_span_action=mc.truncated_span_action, ), columns_to_keep=["messages"], features=MESSAGES_FEATURES, diff --git a/tests/test_row_filters.py b/tests/test_row_filters.py index 666ba88..284b156 100644 --- a/tests/test_row_filters.py +++ b/tests/test_row_filters.py @@ -23,6 +23,12 @@ `tests/test_chat_template_masking.py` pins that span-level behaviour itself. Here it is the input to the filter. +`max_length` splits the outcomes three ways, and the tests below pin each: a +span entirely inside the window is kept intact, a span entirely past it is +dropped, and a span that STRADDLES the window is kept but damaged — it trains +on an answer cut mid-sentence. Only the third is invisible to a keep/drop +count, so it is the one that needs a warning rather than a filter. + No tokenizer is downloaded. `_CharTokenizer` renders the real registered templates through transformers' own generation tracker — the same code path `return_assistant_tokens_mask=True` uses — and then treats one character as @@ -231,6 +237,157 @@ def test_keeps_a_row_whose_assistant_turn_fits_inside_max_length(tokenizer) -> N assert len(kept) == 1 +def _span_bounds(tokenizer, conversation: list[dict]) -> tuple[int, int]: + """(first, last) index of the supervised span. One char == one token here.""" + mask = tokenizer.apply_chat_template( + conversation, return_dict=True, return_assistant_tokens_mask=True + )["assistant_masks"] + return mask.index(1), len(mask) - 1 - mask[::-1].index(1) + + +def test_keeps_a_row_whose_supervised_span_max_length_straddles(tokenizer, caplog) -> None: + """The case a keep/drop count cannot show. + + The span starts inside the window and continues past it, so the row DOES + produce gradient and is kept — while teaching an answer that stops + mid-sentence with no end-of-turn token. Under final-turn-only supervision + the assistant content sits at the end of the row, so this is what a cap + slightly below the row length produces: not a dropped row, a damaged one. + """ + row = [{"role": "user", "content": "q"}, {"role": "assistant", "content": "y" * 200}] + start, end = _span_bounds(tokenizer, row) + straddling = start + 10 + assert start < straddling < end + + with caplog.at_level(logging.INFO, logger=SFT_LOGGER): + kept = _filter_sft_rows( + _sft_dataset([row]), num_proc=1, tokenizer=tokenizer, max_length=straddling + ) + + assert len(kept) == 1 + cut = [record for record in caplog.records if "PART-WAY" in record.getMessage()] + assert [record.levelno for record in cut] == [logging.WARNING] + + +def test_does_not_warn_when_every_supervised_span_fits(tokenizer, caplog) -> None: + row = _exchange(0) + _, end = _span_bounds(tokenizer, row) + + with caplog.at_level(logging.INFO, logger=SFT_LOGGER): + _filter_sft_rows(_sft_dataset([row]), num_proc=1, tokenizer=tokenizer, max_length=end + 1) + + assert not [record for record in caplog.records if "PART-WAY" in record.getMessage()] + + +def test_does_not_warn_about_truncation_when_no_max_length_is_set(tokenizer, caplog) -> None: + """With no cap there is nothing to cut, however long the row is.""" + row = [{"role": "user", "content": "q"}, {"role": "assistant", "content": "y" * 5000}] + + with caplog.at_level(logging.INFO, logger=SFT_LOGGER): + kept = _filter_sft_rows(_sft_dataset([row]), num_proc=1, tokenizer=tokenizer) + + assert len(kept) == 1 + assert not [record for record in caplog.records if "PART-WAY" in record.getMessage()] + + +def test_separates_the_two_reasons_a_row_was_dropped(tokenizer, caplog) -> None: + """The two causes call for opposite fixes — change the data, or raise the + cap — so a single drop count sends the reader to the wrong one half the time. + """ + prompt_only = [{"role": "user", "content": "orphan"}] + beyond_cap = [{"role": "user", "content": "x" * 400}, {"role": "assistant", "content": "a"}] + + with caplog.at_level(logging.INFO, logger=SFT_LOGGER): + _filter_sft_rows( + _sft_dataset([_exchange(0), prompt_only, beyond_cap]), + num_proc=1, + tokenizer=tokenizer, + max_length=64, + ) + + causes = next(record for record in caplog.records if "by cause" in record.getMessage()) + assert "1 with no assistant tokens at all" in causes.getMessage() + assert "1 with every assistant token beyond max_length=64" in causes.getMessage() + + +def _straddling_row(tokenizer) -> tuple[list[dict], int]: + """A row plus a max_length that cuts its supervised span part-way. + + The prompt is long on purpose, so the cut lands deep enough into the row that + an ordinary short row still fits inside the same cap and can serve as the + control. With a short prompt the cap is ~50 tokens, every row exceeds it, and + the filter empties the dataset instead of dropping one row. + """ + row = [{"role": "user", "content": "x" * 300}, {"role": "assistant", "content": "y" * 200}] + start, end = _span_bounds(tokenizer, row) + cut = start + 50 + assert start < cut < end + assert cut > len(tokenizer.render(_exchange(0))), "control row must fit under the cap" + return row, cut + + +def test_a_straddling_row_is_kept_by_default(tokenizer) -> None: + """Default is warn: the row stays, because the fix is normally to raise the + cap rather than to throw the data away.""" + row, cut = _straddling_row(tokenizer) + + kept = _filter_sft_rows(_sft_dataset([row]), num_proc=1, tokenizer=tokenizer, max_length=cut) + + assert len(kept) == 1 + + +def test_a_straddling_row_is_dropped_when_asked(tokenizer) -> None: + """`drop` is for a sequence length fixed by memory, where shortening the + data is the only lever left.""" + row, cut = _straddling_row(tokenizer) + + kept = _filter_sft_rows( + _sft_dataset([row, _exchange(0)]), + num_proc=1, + tokenizer=tokenizer, + max_length=cut, + truncated_span_action="drop", + ) + + assert kept["messages"] == [_exchange(0)] + + +def test_dropping_still_warns(tokenizer, caplog) -> None: + """Drop is not the quiet option. `data.datasets[].weight` multiplies the rows + that survive filtering, so an uneven drop changes the realised mixture — that + has to be visible in the log, not inferred from a row count. + """ + row, cut = _straddling_row(tokenizer) + + with caplog.at_level(logging.INFO, logger=SFT_LOGGER): + _filter_sft_rows( + _sft_dataset([row, _exchange(0)]), + num_proc=1, + tokenizer=tokenizer, + max_length=cut, + truncated_span_action="drop", + ) + + warned = [record for record in caplog.records if "PART-WAY" in record.getMessage()] + assert [record.levelno for record in warned] == [logging.WARNING] + assert "DROPPED" in warned[0].getMessage() + assert "mixture" in warned[0].getMessage() + + +def test_rejects_an_unknown_truncated_span_action(tokenizer) -> None: + """Caught before the dataset is tokenized, not after — a typo must not cost + a full pass over the pool before it surfaces. + """ + with pytest.raises(ValueError, match="truncated_span_action"): + _filter_sft_rows( + _sft_dataset([_exchange(0)]), + num_proc=1, + tokenizer=tokenizer, + max_length=4096, + truncated_span_action="Drop", + ) + + # ── SFT: the failure modes that must be loud ─────────────────────────── @@ -280,6 +437,26 @@ def test_total_drop_error_names_max_length_only_when_it_is_set(tokenizer) -> Non assert "max_length of 4096" in str(with_length.value) +def test_total_drop_error_blames_the_cap_when_every_row_has_an_assistant_turn( + tokenizer, +) -> None: + """Both causes empty the dataset, and they need opposite fixes. When every + row does have an assistant turn, the template is exonerated and saying + otherwise sends the reader on a long detour through their jinja. + """ + rows = [ + [{"role": "user", "content": "x" * 400}, {"role": "assistant", "content": f"a{i}"}] + for i in range(3) + ] + + with pytest.raises(ValueError) as excinfo: + _filter_sft_rows(_sft_dataset(rows), num_proc=1, tokenizer=tokenizer, max_length=32) + + message = str(excinfo.value) + assert "max_seq_length problem" in message + assert "not a template problem" in message + + def test_requires_the_tokenizer_as_a_keyword_argument() -> None: """`build_sft_trainer` binds the tokenizer with `functools.partial`. A missing tokenizer must fail at the call, not as an `AttributeError` inside