From 01c8d03d23c25c47c69b8ee18a105a31b16346c8 Mon Sep 17 00:00:00 2001 From: Konstantin Nikolaou Date: Mon, 24 Aug 2026 15:02:17 +0200 Subject: [PATCH 1/2] fix: warn when max_length cuts a supervised span part-way MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `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. --- src/post_training/methods/sft.py | 160 ++++++++++++++++++++++++------- tests/test_row_filters.py | 99 +++++++++++++++++++ 2 files changed, 225 insertions(+), 34 deletions(-) diff --git a/src/post_training/methods/sft.py b/src/post_training/methods/sft.py index 2f3bc15..97c1016 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,60 @@ # 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 + + +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( @@ -50,13 +105,20 @@ def _filter_sft_rows( tokenizer: PreTrainedTokenizerBase, max_length: int | None = None, ) -> 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 ---------- @@ -92,48 +154,50 @@ 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) 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 (_KEEP, _CUT) 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 +205,34 @@ def _filter_sft_rows( else: logger.info(summary) + # The rows that stay but are damaged. This cannot be inferred from the drop + # count — these rows are NOT dropped, they train on a truncated answer — and + # nothing downstream reports it, so it is a warning at any rate above zero. + if counts[_CUT]: + logger.warning( + "%d of %d rows (%.3f%%) are cut by max_length=%d PART-WAY THROUGH their " + "supervised span. Those rows stay in training, and each one teaches an " + "answer that stops mid-sentence with no end-of-turn token — for reasoning " + "data, an unterminated block. Raise sft.max_seq_length above the " + "rows' length, or drop the over-long rows upstream. Dropping them is safe; " + "training on them is not.", + counts[_CUT], + len(ds), + 100 * counts[_CUT] / len(ds), + max_length, + ) + + 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.", + counts[_NO_ASSISTANT], + counts[_BEYOND_CAP], + max_length, + ) + if dropped: patterns = { " -> ".join(message["role"] for message in row["messages"][-3:]) diff --git a/tests/test_row_filters.py b/tests/test_row_filters.py index 666ba88..80ace0f 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,79 @@ 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() + + # ── SFT: the failure modes that must be loud ─────────────────────────── @@ -280,6 +359,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 From 1320e2e52c6b8234d206cee2e417a89a3310f772 Mon Sep 17 00:00:00 2001 From: Konstantin Nikolaou Date: Tue, 25 Aug 2026 10:25:10 +0200 Subject: [PATCH 2/2] feat: add sft.truncated_span_action to drop rows a cut span damages MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- configs/trl/sft.yaml | 5 ++ src/post_training/config.py | 7 +++ src/post_training/methods/sft.py | 51 ++++++++++++++++----- tests/test_row_filters.py | 78 ++++++++++++++++++++++++++++++++ 4 files changed, 130 insertions(+), 11 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 97c1016..9fe5e59 100644 --- a/src/post_training/methods/sft.py +++ b/src/post_training/methods/sft.py @@ -38,6 +38,9 @@ _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], @@ -104,6 +107,7 @@ 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, and report what does. @@ -133,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 ------- @@ -142,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" ) @@ -163,10 +181,12 @@ def _filter_sft_rows( )["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 (_KEEP, _CUT) else dropped).append(index) + (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 @@ -205,21 +225,28 @@ def _filter_sft_rows( else: logger.info(summary) - # The rows that stay but are damaged. This cannot be inferred from the drop - # count — these rows are NOT dropped, they train on a truncated answer — and - # nothing downstream reports it, so it is a warning at any rate above zero. + # 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. Those rows stay in training, and each one teaches an " - "answer that stops mid-sentence with no end-of-turn token — for reasoning " - "data, an unterminated block. Raise sft.max_seq_length above the " - "rows' length, or drop the over-long rows upstream. Dropping them is safe; " - "training on them is not.", + "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]): @@ -227,10 +254,11 @@ def _filter_sft_rows( # 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.", + "%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: @@ -309,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 80ace0f..284b156 100644 --- a/tests/test_row_filters.py +++ b/tests/test_row_filters.py @@ -310,6 +310,84 @@ def test_separates_the_two_reasons_a_row_was_dropped(tokenizer, caplog) -> None: 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 ───────────────────────────