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