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/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/chat_templates/registry.py b/src/post_training/chat_templates/registry.py index 84b0700..c067e2d 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", @@ -53,6 +62,7 @@ "olmo3-think-sft": "olmo3-think-sft.jinja", "apertus": "apertus.jinja", "tulu3": "tulu3.jinja", + "qwen3": "qwen3.jinja", } diff --git a/src/post_training/chat_templates/templates/qwen3.jinja b/src/post_training/chat_templates/templates/qwen3.jinja new file mode 100644 index 0000000..a8ff3aa --- /dev/null +++ b/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 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 %} + {%- 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') }} + {%- 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 %} + {{- '\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|>' }} + {%- 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\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/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/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..9fe5e59 100644 --- a/src/post_training/methods/sft.py +++ b/src/post_training/methods/sft.py @@ -3,11 +3,14 @@ from __future__ import annotations import logging +from collections import Counter +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 +29,66 @@ logger = logging.getLogger(__name__) +# 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( @@ -38,9 +101,179 @@ ) -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, + truncated_span_action: str = "warn", +) -> Dataset: + """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 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 + ---------- + 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. + 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 + ------- + datasets.Dataset + The rows that stay. + + Raises + ------ + ValueError + 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" + ) + + if len(ds) == 0: + raise ValueError( + "No rows remain after dropping empty 'messages' lists. " + "Check the dataset and its transform." + ) + + # 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", + )["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, 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{scope} " + f"({dropped_percent:.3f}%)." + ) + + if not keep: + 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: + logger.warning(summary) + 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:]) + 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 +325,20 @@ 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, + truncated_span_action=mc.truncated_span_action, + ), columns_to_keep=["messages"], features=MESSAGES_FEATURES, ) 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_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") 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 diff --git a/tests/test_row_filters.py b/tests/test_row_filters.py new file mode 100644 index 0000000..284b156 --- /dev/null +++ b/tests/test_row_filters.py @@ -0,0 +1,562 @@ +"""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. + +`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 +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 + + +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 ─────────────────────────── + + +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_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 + 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