diff --git a/changelog.d/0.77.0/1196.fixed.md b/changelog.d/0.77.0/1196.fixed.md new file mode 100644 index 000000000..0d38c02ce --- /dev/null +++ b/changelog.d/0.77.0/1196.fixed.md @@ -0,0 +1 @@ +- Restored `data.max_length` enforcement for SimPO and KTO on TRL releases that removed `max_prompt_length`, rebuilding the KTO/BCO completion from the untruncated `answer_input_ids` so an over-length prompt no longer discards it, and covering the same measured over-length path for IPO and BCO. For SimPO, KTO, IPO and BCO a prompt is cut only when the row does not fit in `data.max_length`, which replaces the unconditional cut that #2116 introduced (#1185 by @ousamabenyounes in #1196). diff --git a/docs/backends-and-ops.md b/docs/backends-and-ops.md index 29c9e31c7..ad7b12532 100644 --- a/docs/backends-and-ops.md +++ b/docs/backends-and-ops.md @@ -278,10 +278,12 @@ That shows the objective moving, not alignment quality. transformers Trainer does for TRL's CPO. - **Tokens** are ORPO's above: TRL's CPO and ORPO trainers tokenize rows with the same code, so rows that fit in `data.max_length` train on the same tokens - on either backend. Over-long rows are cut with Soup's cap (prompt to - `max_length // 2`, keeping its end, answers to the remaining room). The - transformers SimPO path applies no prompt cap: TRL 0.29.1 cuts only the - answers there, and Soup refuses a row whose shorter answer that cut empties. + on either backend. MLX SimPO, DPO and ORPO cut over-long rows by the prompt + alone (prompt to `max_length // 2`, keeping its end, answers to the remaining + room). The transformers SimPO path now applies that same cap, but only when a + row does not fit `data.max_length` (TRL 0.29.1 cuts only the answers, so a + long prompt reached the model whole); Soup refuses a row whose shorter answer + that cut empties. - **Ties train**, as in TRL: the preference term is the constant `softplus(simpo_gamma)` and the NLL still learns. - **Stops on a nonfinite loss** instead of training on NaN. @@ -375,8 +377,9 @@ training: completion located even when the tokenizer merges the last prompt token into it, BOS prepended to the prompt and EOS appended to the completion when missing. A row longer than `data.max_length` loses the end of its completion; - the prompt is never cut (TRL 0.29.1 has no prompt cap, and the transformers - path adds none). A KL row is re-cut to its own prompt's room. + the prompt is never cut (TRL 0.29.1 has no prompt cap; the transformers path + now caps the prompt only when a row does not fit `data.max_length`). A KL row + is re-cut to its own prompt's room. - **Labels** are read by the loader's `format: kto` converter, exactly as on the transformers backend (`true` / `false`, and the strings and numbers that converter accepts); MLX adds no reading of its own. Only a row handed to the diff --git a/docs/training.md b/docs/training.md index 623231c9d..bfb35f74e 100644 --- a/docs/training.md +++ b/docs/training.md @@ -1361,6 +1361,12 @@ frozen reference. Same input format as DPO; rows are split internally to TRL's BCO unpaired schema (`{prompt, completion, label}`). +When TRL no longer exposes `max_prompt_length`, Soup still enforces +`data.max_length`: BCO, IPO, KTO, and SimPO cut prompts only when a row +would exceed `max_length`. BCO, KTO, and SimPO keep the prompt end; IPO +follows DPO's `truncation_mode` (`keep_start` by default). DPO/ORPO keep +their legacy `max_length // 2` prompt cap. + ```yaml task: bco data: @@ -1935,14 +1941,12 @@ training: quantization: 4bit ``` -**Long prompts are capped as in DPO and ORPO.** trl's CPO trainer cuts only the -answers, so after it tokenizes, `soup train` cuts the prompt to -`data.max_length // 2` tokens, keeping its end, and the answers to the room -left, for the train and validation splits. Every sequence is then at most -`data.max_length`. The cut goes by the prompt alone: a prompt longer than -`data.max_length // 2` tokens loses its start even when prompt and answers -would fit in `data.max_length`. A row whose prompt is at most -`data.max_length // 2` tokens and which fits is unchanged. +**Long prompts are capped only when a row overflows.** trl's CPO trainer cuts +only the answers, so a long prompt reached the model whole. A SimPO prompt is +now cut only when the row does not fit in `data.max_length`: it is then cut to +`data.max_length // 2` tokens, keeping its end, and the answers get the room +left, for the train and validation splits, so every sequence is at most +`data.max_length`. A row that fits is unchanged, whatever its prompt length. **Rows whose completion trl truncates away are refused.** trl's CPO trainer truncates each answer to `data.max_length` minus the **longer** answer's length, diff --git a/src/soup_cli/trainer/_trl_compat.py b/src/soup_cli/trainer/_trl_compat.py index 3294585ec..049fd2d99 100644 --- a/src/soup_cli/trainer/_trl_compat.py +++ b/src/soup_cli/trainer/_trl_compat.py @@ -39,6 +39,20 @@ import inspect from typing import Any +# TRL's own KTO/BCO/CPO/ORPO tokenizers historically truncate an over-length +# prompt with ``keep_end`` — the tail carries a chat template's generation +# header (the assistant turn the completion must follow), so cutting from the +# start would strip that header. KTOConfig exposes no ``truncation_mode`` field +# for the wrappers to read, so the cap KTO applies through +# ``enforce_preference_sequence_limit`` uses this default; it must match TRL's +# historical direction rather than the ``keep_start`` most other configs default +# to. +DEFAULT_PROMPT_TRUNCATION_MODE = "keep_end" + +#: Hugging Face's cross-entropy ignore index; masks prompt tokens out of the +#: completion loss. +_IGNORE_INDEX = -100 + def _installed_trl_version() -> str: """Best-effort version string, for error messages only.""" @@ -145,21 +159,143 @@ def _truncate_tokens(tokens: list[int], limit: int, mode: str) -> list[int]: return tokens[:limit] +def _cap_combined_preference_sequence( + row: dict[str, Any], + prefix: str, + *, + max_length: int, + max_prompt_length: int, + truncation_mode: str, +) -> dict[str, list[int]]: + """Cap a combined prompt/completion row split by the first trainable label.""" + input_ids = list(row[f"{prefix}_input_ids"]) + attention_mask = list(row[f"{prefix}_attention_mask"]) + labels = list(row[f"{prefix}_labels"]) + prompt_size = next( + (index for index, label in enumerate(labels) if label != _IGNORE_INDEX), + len(labels), + ) + prompt_ids = _truncate_tokens( + input_ids[:prompt_size], max_prompt_length, truncation_mode + ) + prompt_mask = _truncate_tokens( + attention_mask[:prompt_size], max_prompt_length, truncation_mode + ) + completion_limit = max(0, max_length - len(prompt_ids)) + completion_ids = input_ids[prompt_size:][:completion_limit] + completion_mask = attention_mask[prompt_size:][:completion_limit] + completion_labels = labels[prompt_size:][:completion_limit] + return { + f"{prefix}_input_ids": prompt_ids + completion_ids, + f"{prefix}_attention_mask": prompt_mask + completion_mask, + f"{prefix}_labels": [_IGNORE_INDEX] * len(prompt_ids) + completion_labels, + } + + +def _rebuild_unpaired_sequence( + prompt_ids: list[int], + prompt_mask: list[int], + answer_ids: list[int], + answer_mask: list[int], + eos_id: int, + *, + max_length: int, + max_prompt_length: int, + truncation_mode: str, +) -> dict[str, list[int]]: + """Build a capped KTO/BCO sequence while keeping EOS as the final token.""" + capped_prompt_ids = _truncate_tokens(list(prompt_ids), max_prompt_length, truncation_mode) + capped_prompt_mask = _truncate_tokens(list(prompt_mask), max_prompt_length, truncation_mode) + answer_ids = list(answer_ids) + answer_mask = list(answer_mask) + # TRL appends EOS to a completion only when it is absent, so an answer that + # already ends in its stop token (a valid short completion is ``[good, ]``) + # must keep the single one. Split the terminal EOS off the body before + # enforcing the budget and re-appending, or ``[good, ]`` trains on + # ``[good, , ]`` on both the matched and KL targets. + has_terminal_eos = bool(answer_ids) and answer_ids[-1] == eos_id + body_ids = answer_ids[:-1] if has_terminal_eos else answer_ids + body_mask = answer_mask[:-1] if has_terminal_eos else answer_mask + body_limit = max(0, max_length - len(capped_prompt_ids) - 1) + capped_body_ids = body_ids[:body_limit] + capped_body_mask = body_mask[:body_limit] + return { + "input_ids": capped_prompt_ids + capped_body_ids + [eos_id], + "attention_mask": capped_prompt_mask + capped_body_mask + [1], + "labels": [_IGNORE_INDEX] * len(capped_prompt_ids) + capped_body_ids + [eos_id], + } + + +def _rotated_answer_index(index: int, row_count: int, chunk_size: int) -> int: + """Return the answer row TRL rotates into KTO's KL completion for ``index``.""" + chunk = max(1, chunk_size) + chunk_start = (index // chunk) * chunk + chunk_end = min(chunk_start + chunk, row_count) + # TRL rotates ``[a[-1]] + a[:-1]``: each row borrows its predecessor's answer. + return index - 1 if index > chunk_start else chunk_end - 1 + + +def _rebuild_unpaired_completion_row( + row: dict[str, Any], + *, + max_length: int, + max_prompt_length: int, + truncation_mode: str, +) -> dict[str, list[int]]: + """Reassemble the KTO/BCO completion from the untruncated answer tokens. + + TRL's KTO/BCO tokenizer caps the answer to ``max_length - len(prompt)`` + while ignoring ``max_prompt_length`` (``_process_tokens`` truncates + ``answer_input_ids`` before concatenating). An over-length prompt therefore + fills the whole budget and leaves ``completion_input_ids`` with the answer + already sliced away, so capping that column can only recover an empty + completion. The matched answer survives verbatim in ``answer_input_ids`` / + ``answer_attention_mask`` (columns the loss never reads) and the prompt in + ``prompt_input_ids``, so rebuild the combined sequence from those before the + real cap runs. TRL guarantees ``completion_input_ids`` ends with EOS; reuse + that token so a prompt-driven cap never drops the stop signal. + """ + eos_id = list(row["completion_input_ids"])[-1] + rebuilt = _rebuild_unpaired_sequence( + list(row["prompt_input_ids"]), + list(row["prompt_attention_mask"]), + list(row["answer_input_ids"]), + list(row["answer_attention_mask"]), + eos_id, + max_length=max_length, + max_prompt_length=max_prompt_length, + truncation_mode=truncation_mode, + ) + return { + "completion_input_ids": rebuilt["input_ids"], + "completion_attention_mask": rebuilt["attention_mask"], + "completion_labels": rebuilt["labels"], + } + + def enforce_preference_sequence_limit( dataset: Any, *, max_length: int, max_prompt_length: int, truncation_mode: str, + only_when_overflow: bool = False, + kl_chunk_size: int | None = None, ) -> Any: """Restore TRL's removed prompt cap on an already-tokenized dataset. - TRL 0.29 exposes two preference dataset layouts. DPO stores a shared - ``prompt_ids`` plus separate completion ids; experimental ORPO stores two - combined sequences whose prompt span is identified by ``-100`` labels. - Applying the cap after TRL tokenizes keeps text and conversational inputs - on the exact same chat-template path while guaranteeing that the tensors - reaching the model obey ``data.max_length``. + TRL 0.29 exposes a few preference dataset layouts. DPO stores a shared + ``prompt_ids`` plus separate completion ids; experimental ORPO/CPO store + two combined sequences whose prompt span is identified by ``-100`` labels; + BCO/KTO store one unpaired ``completion_*`` sequence, with KTO adding a + ``KL_completion_*`` sequence as well. Applying the cap after TRL tokenizes + keeps text and conversational inputs on the exact same chat-template path + while guaranteeing that the tensors reaching the model obey + ``data.max_length``. + + ``only_when_overflow`` preserves rows whose TRL-built model inputs already + fit inside ``max_length``; DPO and ORPO keep the legacy unconditional cap, + while BCO, IPO, KTO and SimPO use the conditional path. """ columns = set(getattr(dataset, "column_names", ())) dpo_columns = {"prompt_ids", "chosen_ids", "rejected_ids"} @@ -173,10 +309,21 @@ def enforce_preference_sequence_limit( "rejected_attention_mask", "rejected_labels", } + unpaired_columns = { + "prompt_input_ids", + "prompt_attention_mask", + "completion_input_ids", + "completion_attention_mask", + "completion_labels", + } if dpo_columns <= columns: def cap_dpo(row: dict[str, Any]) -> dict[str, Any]: + if only_when_overflow and len(row["prompt_ids"]) + max( + len(row["chosen_ids"]), len(row["rejected_ids"]) + ) <= max_length: + return {} prompt = _truncate_tokens( list(row["prompt_ids"]), max_prompt_length, truncation_mode ) @@ -191,45 +338,135 @@ def cap_dpo(row: dict[str, Any]) -> dict[str, Any]: if orpo_columns <= columns: - def cap_combined(row: dict[str, Any], prefix: str) -> dict[str, list[int]]: - input_ids = list(row[f"{prefix}_input_ids"]) - attention_mask = list(row[f"{prefix}_attention_mask"]) - labels = list(row[f"{prefix}_labels"]) - prompt_size = next( - (index for index, label in enumerate(labels) if label != -100), - len(labels), - ) + def cap_orpo(row: dict[str, Any]) -> dict[str, Any]: + if only_when_overflow and ( + len(row["chosen_input_ids"]) <= max_length + and len(row["rejected_input_ids"]) <= max_length + ): + return {} prompt_ids = _truncate_tokens( - input_ids[:prompt_size], max_prompt_length, truncation_mode + list(row["prompt_input_ids"]), max_prompt_length, truncation_mode ) prompt_mask = _truncate_tokens( - attention_mask[:prompt_size], max_prompt_length, truncation_mode + list(row["prompt_attention_mask"]), max_prompt_length, truncation_mode ) - completion_limit = max(0, max_length - len(prompt_ids)) - completion_ids = input_ids[prompt_size:][:completion_limit] - completion_mask = attention_mask[prompt_size:][:completion_limit] - completion_labels = labels[prompt_size:][:completion_limit] return { - f"{prefix}_input_ids": prompt_ids + completion_ids, - f"{prefix}_attention_mask": prompt_mask + completion_mask, - f"{prefix}_labels": [-100] * len(prompt_ids) + completion_labels, + "prompt_input_ids": prompt_ids, + "prompt_attention_mask": prompt_mask, + **_cap_combined_preference_sequence( + row, + "chosen", + max_length=max_length, + max_prompt_length=max_prompt_length, + truncation_mode=truncation_mode, + ), + **_cap_combined_preference_sequence( + row, + "rejected", + max_length=max_length, + max_prompt_length=max_prompt_length, + truncation_mode=truncation_mode, + ), } - def cap_orpo(row: dict[str, Any]) -> dict[str, Any]: + return dataset.map(cap_orpo) + + if unpaired_columns <= columns: + answer_columns = {"answer_input_ids", "answer_attention_mask"} + has_answer = answer_columns <= columns + kl_columns = { + "KL_prompt_input_ids", + "KL_prompt_attention_mask", + "KL_completion_input_ids", + "KL_completion_attention_mask", + "KL_completion_labels", + } + has_kl = kl_columns <= columns + answers = list(dataset["answer_input_ids"]) if has_answer and has_kl else [] + answer_masks = list(dataset["answer_attention_mask"]) if answers else [] + kl_chunk = max(1, int(kl_chunk_size or 1)) + + def cap_unpaired(row: dict[str, Any], index: int) -> dict[str, Any]: + fits = len(row["completion_input_ids"]) <= max_length and ( + not has_kl or len(row["KL_completion_input_ids"]) <= max_length + ) + if only_when_overflow and fits: + return {} prompt_ids = _truncate_tokens( list(row["prompt_input_ids"]), max_prompt_length, truncation_mode ) prompt_mask = _truncate_tokens( list(row["prompt_attention_mask"]), max_prompt_length, truncation_mode ) - return { + # TRL already sliced the answer out of ``completion_input_ids`` when + # the prompt overflowed; rebuild it from the untruncated answer + # columns so the cap keeps real completion tokens, not just the EOS. + completion_source = ( + _rebuild_unpaired_completion_row( + row, + max_length=max_length, + max_prompt_length=max_prompt_length, + truncation_mode=truncation_mode, + ) + if has_answer + else row + ) + capped = { "prompt_input_ids": prompt_ids, "prompt_attention_mask": prompt_mask, - **cap_combined(row, "chosen"), - **cap_combined(row, "rejected"), + **_cap_combined_preference_sequence( + completion_source, + "completion", + max_length=max_length, + max_prompt_length=max_prompt_length, + truncation_mode=truncation_mode, + ), } - - return dataset.map(cap_orpo) + if {"KL_prompt_input_ids", "KL_prompt_attention_mask"} <= columns: + capped["KL_prompt_input_ids"] = _truncate_tokens( + list(row["KL_prompt_input_ids"]), + max_prompt_length, + truncation_mode, + ) + capped["KL_prompt_attention_mask"] = _truncate_tokens( + list(row["KL_prompt_attention_mask"]), + max_prompt_length, + truncation_mode, + ) + if { + "KL_completion_input_ids", + "KL_completion_attention_mask", + "KL_completion_labels", + } <= columns: + if answers: + eos_id = list(row["KL_completion_input_ids"])[-1] + lender = _rotated_answer_index(index, len(answers), kl_chunk) + kl = _rebuild_unpaired_sequence( + list(row["KL_prompt_input_ids"]), + list(row["KL_prompt_attention_mask"]), + list(answers[lender]), + list(answer_masks[lender]), + eos_id, + max_length=max_length, + max_prompt_length=max_prompt_length, + truncation_mode=truncation_mode, + ) + capped["KL_completion_input_ids"] = kl["input_ids"] + capped["KL_completion_attention_mask"] = kl["attention_mask"] + capped["KL_completion_labels"] = kl["labels"] + else: + capped.update( + _cap_combined_preference_sequence( + row, + "KL_completion", + max_length=max_length, + max_prompt_length=max_prompt_length, + truncation_mode=truncation_mode, + ) + ) + return capped + + return dataset.map(cap_unpaired, with_indices=True) raise ValueError( "TRL prepared a preference dataset with an unknown token layout; " diff --git a/src/soup_cli/trainer/bco.py b/src/soup_cli/trainer/bco.py index 96535e26e..143ec724c 100644 --- a/src/soup_cli/trainer/bco.py +++ b/src/soup_cli/trainer/bco.py @@ -112,6 +112,8 @@ def setup(self, dataset: dict) -> None: from datasets import Dataset from soup_cli.trainer._trl_compat import ( + config_accepts, + enforce_preference_sequence_limit, prompt_length_kwargs, resolve_trl_symbol, ) @@ -238,6 +240,20 @@ def setup(self, dataset: dict) -> None: eval_dataset=eval_ds, processing_class=self.tokenizer, ) + if not config_accepts(bco_config_cls, "max_prompt_length"): + cap_kwargs = { + "max_length": cfg.data.max_length, + "max_prompt_length": cfg.data.max_length // 2, + "truncation_mode": bco_config.truncation_mode, + "only_when_overflow": True, + } + self.trainer.train_dataset = enforce_preference_sequence_limit( + self.trainer.train_dataset, **cap_kwargs + ) + if self.trainer.eval_dataset is not None: + self.trainer.eval_dataset = enforce_preference_sequence_limit( + self.trainer.eval_dataset, **cap_kwargs + ) # #359 - the same exposure #336 fixed in sft.py: with LoRA the # no-decay optimizer group is empty, DeepSpeed drops it, and the LR diff --git a/src/soup_cli/trainer/ipo.py b/src/soup_cli/trainer/ipo.py index 37f99c8d6..89f5e1cd7 100644 --- a/src/soup_cli/trainer/ipo.py +++ b/src/soup_cli/trainer/ipo.py @@ -78,7 +78,11 @@ def setup(self, dataset: dict) -> None: from datasets import Dataset from trl import DPOConfig, DPOTrainer - from soup_cli.trainer._trl_compat import prompt_length_kwargs + from soup_cli.trainer._trl_compat import ( + config_accepts, + enforce_preference_sequence_limit, + prompt_length_kwargs, + ) from soup_cli.trainer.sft import _enable_hf_transfer_progress _enable_hf_transfer_progress() @@ -188,6 +192,20 @@ def setup(self, dataset: dict) -> None: eval_dataset=eval_ds, processing_class=self.tokenizer, ) + if not config_accepts(DPOConfig, "max_prompt_length"): + cap_kwargs = { + "max_length": cfg.data.max_length, + "max_prompt_length": cfg.data.max_length // 2, + "truncation_mode": dpo_config.truncation_mode, + "only_when_overflow": True, + } + self.trainer.train_dataset = enforce_preference_sequence_limit( + self.trainer.train_dataset, **cap_kwargs + ) + if self.trainer.eval_dataset is not None: + self.trainer.eval_dataset = enforce_preference_sequence_limit( + self.trainer.eval_dataset, **cap_kwargs + ) # #359 - the same exposure #336 fixed in sft.py: with LoRA the # no-decay optimizer group is empty, DeepSpeed drops it, and the LR diff --git a/src/soup_cli/trainer/kto.py b/src/soup_cli/trainer/kto.py index c129f465c..7ba96676c 100644 --- a/src/soup_cli/trainer/kto.py +++ b/src/soup_cli/trainer/kto.py @@ -78,7 +78,12 @@ def setup(self, dataset: dict) -> None: from datasets import Dataset from trl import KTOConfig, KTOTrainer - from soup_cli.trainer._trl_compat import prompt_length_kwargs + from soup_cli.trainer._trl_compat import ( + DEFAULT_PROMPT_TRUNCATION_MODE, + config_accepts, + enforce_preference_sequence_limit, + prompt_length_kwargs, + ) # Enable Rich progress bar for HuggingFace downloads from soup_cli.trainer.sft import _enable_hf_transfer_progress @@ -217,6 +222,21 @@ def setup(self, dataset: dict) -> None: eval_dataset=eval_ds, processing_class=self.tokenizer, ) + if not config_accepts(KTOConfig, "max_prompt_length"): + cap_kwargs = { + "max_length": cfg.data.max_length, + "max_prompt_length": cfg.data.max_length // 2, + "truncation_mode": DEFAULT_PROMPT_TRUNCATION_MODE, + "only_when_overflow": True, + "kl_chunk_size": batch_size, + } + self.trainer.train_dataset = enforce_preference_sequence_limit( + self.trainer.train_dataset, **cap_kwargs + ) + if self.trainer.eval_dataset is not None: + self.trainer.eval_dataset = enforce_preference_sequence_limit( + self.trainer.eval_dataset, **cap_kwargs + ) # #359 - the same exposure #336 fixed in sft.py: with LoRA the # no-decay optimizer group is empty, DeepSpeed drops it, and the LR diff --git a/src/soup_cli/trainer/simpo.py b/src/soup_cli/trainer/simpo.py index 94cb024f2..3cf778cb7 100644 --- a/src/soup_cli/trainer/simpo.py +++ b/src/soup_cli/trainer/simpo.py @@ -233,6 +233,7 @@ def setup(self, dataset: dict) -> None: "max_length": cfg.data.max_length, "max_prompt_length": cfg.data.max_length // 2, "truncation_mode": cpo_config.truncation_mode, + "only_when_overflow": True, } self.trainer.train_dataset = enforce_preference_sequence_limit( self.trainer.train_dataset, **cap_kwargs diff --git a/tests/test_mlx_kto.py b/tests/test_mlx_kto.py index 642a17435..6c29ec500 100644 --- a/tests/test_mlx_kto.py +++ b/tests/test_mlx_kto.py @@ -93,7 +93,7 @@ def test_the_effective_settings_are_ktoconfig_defaults(self): def test_the_transformers_kto_path_passes_beta_and_max_length_only(self): """Soup builds KTOConfig with kto_beta and max_length; the prompt-length - helper forwards nothing on 0.29.1, and no Soup cap touches the rows.""" + helper forwards nothing on 0.29.1.""" pytest.importorskip("trl") from trl.experimental.kto import KTOConfig @@ -111,7 +111,6 @@ def test_the_transformers_kto_path_passes_beta_and_max_length_only(self): assert not names & {"loss_type", "desirable_weight", "undesirable_weight", "precompute_ref_log_probs", "disable_dropout"} assert prompt_length_kwargs(KTOConfig, 64) == {} - assert "enforce_preference_sequence_limit" not in source def test_the_training_sampler_is_sequential(self): """``_get_train_sampler`` (kto_trainer.py:1315-1320): TRL's KTO never diff --git a/tests/test_simpo_prompt_cap.py b/tests/test_simpo_prompt_cap.py index 32a6b3286..8250015ac 100644 --- a/tests/test_simpo_prompt_cap.py +++ b/tests/test_simpo_prompt_cap.py @@ -168,9 +168,10 @@ def _prompt_words_for(tmp_path, monkeypatch, n_tokens: int) -> int: class TestThePromptBudgetBoundary: - """The cap cuts by the prompt alone: any prompt over ``max_length // 2`` - tokens loses its start, even when prompt plus answers fit in - ``max_length``; one of exactly ``max_length // 2`` tokens is untouched.""" + """The cap cuts only when the row overflows ``max_length``: a prompt over + ``max_length // 2`` tokens is left whole while the whole row fits, and loses + its start down to its last ``max_length // 2`` tokens once the row + overflows; one of exactly ``max_length // 2`` tokens is untouched.""" def test_a_prompt_of_exactly_the_budget_is_unchanged(self, tmp_path, monkeypatch): raw = _row(_prompt_words_for(tmp_path, monkeypatch, PROMPT_BUDGET), 2, 1) @@ -181,25 +182,43 @@ def test_a_prompt_of_exactly_the_budget_is_unchanged(self, tmp_path, monkeypatch for key, value in before.items(): assert after[key] == value, key - def test_one_token_over_the_budget_is_cut_although_the_row_fits( + def test_one_token_over_the_budget_is_untouched_while_the_row_fits( self, tmp_path, monkeypatch ): raw = _row(_prompt_words_for(tmp_path, monkeypatch, PROMPT_BUDGET + 1), 2, 1) wrapper = _setup_simpo(tmp_path, monkeypatch, [raw] * 4) before = _before(wrapper, raw) assert len(before["prompt_input_ids"]) == PROMPT_BUDGET + 1 - # The whole row already fits: the cut comes from the prompt budget alone. + # The whole row fits, so the over-budget prompt is left whole. assert all(len(before[key]) <= MAX_LENGTH for key in SEQUENCE_KEYS), before + after = wrapper.trainer.train_dataset[0] + for key, value in before.items(): + assert after[key] == value, key + + def test_an_over_budget_prompt_is_cut_to_the_budget_once_the_row_overflows( + self, tmp_path, monkeypatch + ): + # TRL's CPO tokenizer already fits any row whose prompt leaves room for + # its answers, so a prompt reaches the cap only once the row it builds + # overflows ``max_length``. A prompt one token over ``max_length`` is the + # tightest such row; its cut must land on exactly ``PROMPT_BUDGET``. + raw = _row(_prompt_words_for(tmp_path, monkeypatch, MAX_LENGTH + 1), 2, 1) + wrapper = _setup_simpo(tmp_path, monkeypatch, [raw] * 4) + before = _before(wrapper, raw) + assert len(before["prompt_input_ids"]) == MAX_LENGTH + 1 + # The row does not fit: the cap trims the prompt to its last PROMPT_BUDGET. + assert any(len(before[key]) > MAX_LENGTH for key in SEQUENCE_KEYS), before + after = wrapper.trainer.train_dataset[0] kept = before["prompt_input_ids"][-PROMPT_BUDGET:] + assert len(after["prompt_input_ids"]) == PROMPT_BUDGET assert after["prompt_input_ids"] == kept assert after["prompt_attention_mask"] == [1] * PROMPT_BUDGET for side in ("chosen", "rejected"): assert after[f"{side}_input_ids"][:PROMPT_BUDGET] == kept, side assert after[f"{side}_labels"][:PROMPT_BUDGET] == [-100] * PROMPT_BUDGET, side - assert len(after[f"{side}_input_ids"]) == len(before[f"{side}_input_ids"]) - 1 - assert _trainable(after, side) == _trainable(before, side), side + assert len(after[f"{side}_input_ids"]) <= MAX_LENGTH, side class TestTheCapAndThe1208Refusal: diff --git a/tests/test_trl_preference_config_contract.py b/tests/test_trl_preference_config_contract.py index 5f1e78c17..e624c971c 100644 --- a/tests/test_trl_preference_config_contract.py +++ b/tests/test_trl_preference_config_contract.py @@ -79,8 +79,13 @@ def _pref_rows(n=8): return [{"prompt": "hi", "chosen": " good answer", "rejected": " bad"} for _ in range(n)] +_MAX_SEQUENCE_LENGTH = 64 +_MAX_PROMPT_LENGTH = _MAX_SEQUENCE_LENGTH // 2 +_OVERLONG_PROMPT_WORDS = 161 + + def _long_pref_rows(n=8): - prompt = " ".join(["hello"] * 161) + prompt = " ".join(["hello"] * _OVERLONG_PROMPT_WORDS) return [{"prompt": prompt, "chosen": " good answer", "rejected": " bad"} for _ in range(n)] @@ -88,9 +93,62 @@ def _kto_rows(n=8): return [{"prompt": "hi", "completion": " good answer", "label": i % 2 == 0} for i in range(n)] +def _long_kto_rows(n=8): + prompt = " ".join(["hello"] * _OVERLONG_PROMPT_WORDS) + return [{"prompt": prompt, "completion": " good answer", "label": i % 2 == 0} for i in range(n)] + + +#: Every word here is a single token in the tiny tokenizer, so the completion is +#: this many tokens long — comfortably over ``_MAX_SEQUENCE_LENGTH`` on its own. +_LONG_COMPLETION_WORDS = 100 +_LONG_COMPLETION = " " + " ".join(["answer"] * _LONG_COMPLETION_WORDS) + +#: The label used to mask prompt tokens out of the completion loss. +_IGNORE_LABEL = -100 + +_FALLBACK_MAX_SEQUENCE_LENGTH = 4 +_FALLBACK_MAX_PROMPT_LENGTH = 2 +_FALLBACK_PROMPT_TOKENS = [101, 102, 103, 104] +_FALLBACK_COMPLETION_TOKENS = [201, 202, 203] +_FALLBACK_KL_COMPLETION_TOKENS = [301, 302, 303] +_FALLBACK_ATTENTION_MASK = [1] + + +def _mixed_kto_rows(n=8): + """Half the rows overflow on the prompt, half on the completion. + + The prompt-heavy rows are the regression: TRL slices the answer to fit + ``max_length`` before Soup's cap runs, so a wrapper that re-derives the + completion from TRL's own truncated column keeps only the trailing EOS. The + completion-heavy rows exercise the ordinary answer-truncation path. + """ + long_prompt = " ".join(["hello"] * _OVERLONG_PROMPT_WORDS) + rows = [] + for i in range(n): + if i % 2 == 0: + rows.append({"prompt": long_prompt, "completion": " good answer", "label": True}) + else: + rows.append({"prompt": "hi", "completion": _LONG_COMPLETION, "label": False}) + return rows + + +def _mixed_pref_rows(n=8): + """BCO's preference-shaped analogue of :func:`_mixed_kto_rows`.""" + long_prompt = " ".join(["hello"] * _OVERLONG_PROMPT_WORDS) + rows = [] + for i in range(n): + if i % 2 == 0: + rows.append({"prompt": long_prompt, "chosen": " good answer", "rejected": " bad"}) + else: + rows.append({"prompt": "hi", "chosen": _LONG_COMPLETION, "rejected": " bad"}) + return rows + + #: KTO refuses `per_device_train_batch_size == 1` outright — its KL term is #: degenerate at 1 — so the floor is per task, not global. _MIN_BATCH = {"kto": 2} +# At batch 2 both KL rotation directions coincide; 4 tells them apart. +_KL_ROTATION_BATCH = 4 _ROWS = { "bco": _pref_rows, @@ -100,6 +158,7 @@ def _kto_rows(n=8): "orpo": _pref_rows, "simpo": _pref_rows, } +_PREPARED_ROW_MULTIPLIER = {"bco": 2} #: task -> (wrapper import path, the trl config class its `setup()` builds). #: The config class is named so a failure says *which* trl API moved, rather @@ -132,7 +191,11 @@ def _cfg(weights, out_dir, task): "task": task, "backend": "transformers", "modality": "text", - "data": {"train": "train.jsonl", "max_length": 64, "chat_template": "chatml"}, + "data": { + "train": "train.jsonl", + "max_length": _MAX_SEQUENCE_LENGTH, + "chat_template": "chatml", + }, "training": { "batch_size": _MIN_BATCH.get(task, 1), "quantization": "none", @@ -191,25 +254,35 @@ def test_the_config_carries_the_arguments_the_wrapper_passes( wrapper = _build(tmp_path, monkeypatch, task) args = wrapper.trainer.args if config_accepts(type(args), "max_prompt_length"): - assert args.max_prompt_length == 32, (task, args.max_prompt_length) + assert args.max_prompt_length == _MAX_PROMPT_LENGTH, (task, args.max_prompt_length) else: assert not hasattr(args, "max_prompt_length"), ( f"{task}: installed trl does not accept `max_prompt_length` as a " f"keyword, yet the built config has the attribute — the capability " f"probe and the object disagree" ) - assert args.max_length == 64, (task, args.max_length) - - @pytest.mark.parametrize("task", ("dpo", "orpo", "simpo")) + assert args.max_length == _MAX_SEQUENCE_LENGTH, (task, args.max_length) + + @pytest.mark.parametrize( + ("task", "rows_factory"), + ( + ("bco", _long_pref_rows), + ("dpo", _long_pref_rows), + ("ipo", _long_pref_rows), + ("kto", _long_kto_rows), + ("orpo", _long_pref_rows), + ("simpo", _long_pref_rows), + ), + ) def test_removed_prompt_cap_is_enforced_on_the_effective_batch( - self, tmp_path, monkeypatch, task + self, tmp_path, monkeypatch, task, rows_factory ): """TRL 0.29 must not turn ``data.max_length`` into a cosmetic field. The assertion is on the tensors the model receives, not the config or an intermediate column: 0.29 accepted ``max_length=64`` while emitting - 164/165-token DPO/ORPO batches from this exact 161-token prompt, and - SimPO did the same until #2112 gave it the cap. + over-length batches from this exact 161-token prompt (164/165-token + DPO/ORPO batches, and SimPO did the same until #2112 gave it the cap). """ from soup_cli.trainer._trl_compat import config_accepts @@ -221,19 +294,78 @@ def test_removed_prompt_cap_is_enforced_on_the_effective_batch( monkeypatch.chdir(tmp_path) cfg = _cfg(weights, tmp_path / "out", task) wrapper = getattr(importlib.import_module(module), cls_name)(cfg, device="cpu") - wrapper.setup({"train": _long_pref_rows(8)}) + wrapper.setup({"train": rows_factory(8)}) if config_accepts(type(wrapper.trainer.args), "max_prompt_length"): pytest.skip("installed TRL still enforces its own prompt cap") - assert len(wrapper.trainer.train_dataset) == 8 + expected_rows = 8 * _PREPARED_ROW_MULTIPLIER.get(task, 1) + assert len(wrapper.trainer.train_dataset) == expected_rows batch = next(iter(wrapper.trainer.get_train_dataloader())) sequence_keys = [key for key in batch if key.endswith("input_ids")] assert sequence_keys, batch.keys() - assert all(batch[key].shape[-1] <= 64 for key in sequence_keys), { + assert all(batch[key].shape[-1] <= _MAX_SEQUENCE_LENGTH for key in sequence_keys), { key: tuple(batch[key].shape) for key in sequence_keys } + @pytest.mark.parametrize( + ("task", "rows_factory"), + ( + ("kto", _mixed_kto_rows), + ("bco", _mixed_pref_rows), + ), + ) + def test_unpaired_completion_content_survives_the_cap( + self, tmp_path, monkeypatch, task, rows_factory + ): + """The cap must keep the *content* of the completion, not just its shape. + + The shape check above passes even when the completion is empty — an + over-length prompt makes TRL slice the answer away inside + ``completion_input_ids`` before Soup's cap runs, so a row can be exactly + ``max_length`` tokens yet train on nothing but a trailing EOS. This + asserts, per row, that the unmasked completion tokens are the real + leading answer tokens (``answer_input_ids``), including at least one + non-EOS token. + + The assertion reads content off ``completion_labels``, never off + ``answer_input_ids`` itself: that column is TRL scratch the loss ignores + and Soup does not cap, so a generic ``*input_ids <= max_length`` sweep + would wrongly flag it on the long-completion rows. + """ + from soup_cli.trainer._trl_compat import config_accepts + + module, cls_name, _ = _WRAPPERS[task] + import importlib + + from transformers import AutoTokenizer + + weights = _tiny_llama_dir(tmp_path) + _write_tiny_tokenizer(weights) + monkeypatch.chdir(tmp_path) + eos_id = AutoTokenizer.from_pretrained(weights).eos_token_id + cfg = _cfg(weights, tmp_path / "out", task) + wrapper = getattr(importlib.import_module(module), cls_name)(cfg, device="cpu") + wrapper.setup({"train": rows_factory(8)}) + + if config_accepts(type(wrapper.trainer.args), "max_prompt_length"): + pytest.skip("installed TRL still enforces its own prompt cap") + + dataset = wrapper.trainer.train_dataset + assert "answer_input_ids" in dataset.column_names, dataset.column_names + for index, row in enumerate(dataset): + completion_ids = list(row["completion_input_ids"]) + labels = list(row["completion_labels"]) + assert len(completion_ids) <= _MAX_SEQUENCE_LENGTH, (task, index, len(completion_ids)) + unmasked = [tok for tok, lab in zip(completion_ids, labels) if lab != _IGNORE_LABEL] + non_eos = [tok for tok in unmasked if tok != eos_id] + assert non_eos, ( + f"{task} row {index}: the capped completion is EOS-only — the answer " + f"was truncated away before the cap, so the row trains on nothing" + ) + answer = list(row["answer_input_ids"]) + assert non_eos == answer[: len(non_eos)], (task, index, non_eos, answer) + class TestTheCanaryCoversWhatItClaims: """Guards the *coverage* property, not the code. The bug shipped because a @@ -307,3 +439,536 @@ def test_the_installed_trl_provides_every_symbol_the_trainers_import(self): f"public/experimental resolver — the [train] bounds and the code " f"disagree: {broken}" ) + + +class TestKtoTruncatesThePromptFromTheEnd: + """KTOConfig exposes no ``truncation_mode`` field, so KTO caps its prompt + with ``DEFAULT_PROMPT_TRUNCATION_MODE``. That must be ``keep_end`` — TRL's + own historical direction for KTO/BCO/CPO/ORPO — because a chat template puts + the generation header (the assistant turn the completion continues) at the + END of the prompt; ``keep_start`` would cut it off. + """ + + def test_the_kto_cap_keeps_the_prompt_tail_not_its_head(self): + from datasets import Dataset + + from soup_cli.trainer._trl_compat import ( + DEFAULT_PROMPT_TRUNCATION_MODE, + enforce_preference_sequence_limit, + ) + + assert DEFAULT_PROMPT_TRUNCATION_MODE == "keep_end" + + eos_id = 999 + prompt = list(range(100, 116)) # 16 distinct, order-revealing prompt tokens + answer = [200, 201] + max_prompt_length = 4 + max_length = 8 + dataset = Dataset.from_dict( + { + "prompt_input_ids": [prompt], + "prompt_attention_mask": [[1] * len(prompt)], + "answer_input_ids": [answer], + "answer_attention_mask": [[1] * len(answer)], + "completion_input_ids": [prompt + answer + [eos_id]], + "completion_attention_mask": [[1] * (len(prompt) + len(answer) + 1)], + "completion_labels": [[_IGNORE_LABEL] * len(prompt) + answer + [eos_id]], + } + ) + + capped = enforce_preference_sequence_limit( + dataset, + max_length=max_length, + max_prompt_length=max_prompt_length, + truncation_mode=DEFAULT_PROMPT_TRUNCATION_MODE, + )[0] + + tail = prompt[-max_prompt_length:] + head = prompt[:max_prompt_length] + assert list(capped["prompt_input_ids"]) == tail + assert list(capped["prompt_input_ids"]) != head + # the completion's masked prompt span must keep the same tail, so the + # generation header survives on the sequence the model is trained on + masked = [ + tok + for tok, label in zip(capped["completion_input_ids"], capped["completion_labels"]) + if label == _IGNORE_LABEL + ] + assert masked == tail + + +class TestUnpairedKlFallbackWithoutAnswerColumns: + def test_kl_completion_caps_without_answer_columns(self): + from datasets import Dataset + + from soup_cli.trainer._trl_compat import enforce_preference_sequence_limit + + prompt = _FALLBACK_PROMPT_TOKENS + completion = _FALLBACK_COMPLETION_TOKENS + kl_completion = _FALLBACK_KL_COMPLETION_TOKENS + prompt_mask = _FALLBACK_ATTENTION_MASK * len(prompt) + completion_mask = _FALLBACK_ATTENTION_MASK * len(completion) + kl_completion_mask = _FALLBACK_ATTENTION_MASK * len(kl_completion) + dataset = Dataset.from_dict( + { + "prompt_input_ids": [prompt], + "prompt_attention_mask": [prompt_mask], + "completion_input_ids": [prompt + completion], + "completion_attention_mask": [prompt_mask + completion_mask], + "completion_labels": [[_IGNORE_LABEL] * len(prompt) + completion], + "KL_prompt_input_ids": [prompt], + "KL_prompt_attention_mask": [prompt_mask], + "KL_completion_input_ids": [prompt + kl_completion], + "KL_completion_attention_mask": [prompt_mask + kl_completion_mask], + "KL_completion_labels": [[_IGNORE_LABEL] * len(prompt) + kl_completion], + } + ) + + capped = enforce_preference_sequence_limit( + dataset, + max_length=_FALLBACK_MAX_SEQUENCE_LENGTH, + max_prompt_length=_FALLBACK_MAX_PROMPT_LENGTH, + truncation_mode="keep_end", + )[0] + + expected_prompt = prompt[-_FALLBACK_MAX_PROMPT_LENGTH:] + expected_kl_completion = kl_completion[ + : _FALLBACK_MAX_SEQUENCE_LENGTH - len(expected_prompt) + ] + assert list(capped["KL_completion_input_ids"]) == ( + expected_prompt + expected_kl_completion + ) + assert list(capped["KL_completion_labels"]) == ( + [_IGNORE_LABEL] * len(expected_prompt) + expected_kl_completion + ) + + +class TestUnpairedRebuildPreservesAnExistingTerminalEos: + """An already-terminal EOS must survive the rebuild, not be doubled. + + When the prompt overflows, ``cap_unpaired`` rebuilds the KTO/BCO completion + (and KTO's KL completion) from the untruncated ``answer_input_ids``. TRL + stores the answer *with* its stop token — a valid short completion is + ``[good, ]`` — and the original TRL path appends EOS only when it is + absent. Appending unconditionally trains the matched and KL targets on + ``[good, , ]``, changing the target for inputs that already carry + their stop token. The oracle is the exact answer the row was built with. + """ + + def test_a_terminal_eos_is_kept_single_on_matched_and_kl_targets(self): + from datasets import Dataset + + from soup_cli.trainer._trl_compat import ( + DEFAULT_PROMPT_TRUNCATION_MODE, + enforce_preference_sequence_limit, + ) + + eos_id = 999 + prompt = list(range(100, 116)) # 16 tokens, overflows the prompt cap + answer = [200, eos_id] # a short, valid completion that already ends in EOS + max_prompt_length = 4 + max_length = 8 + # KTO batches run the KL rotation; pin a chunk so the branch is exercised. + kl_chunk_size = 2 + dataset = Dataset.from_dict( + { + "prompt_input_ids": [prompt], + "prompt_attention_mask": [[1] * len(prompt)], + "answer_input_ids": [answer], + "answer_attention_mask": [[1] * len(answer)], + "completion_input_ids": [prompt + answer], + "completion_attention_mask": [[1] * (len(prompt) + len(answer))], + "completion_labels": [[_IGNORE_LABEL] * len(prompt) + answer], + "KL_prompt_input_ids": [prompt], + "KL_prompt_attention_mask": [[1] * len(prompt)], + "KL_completion_input_ids": [prompt + answer], + "KL_completion_attention_mask": [[1] * (len(prompt) + len(answer))], + "KL_completion_labels": [[_IGNORE_LABEL] * len(prompt) + answer], + } + ) + + capped = enforce_preference_sequence_limit( + dataset, + max_length=max_length, + max_prompt_length=max_prompt_length, + truncation_mode=DEFAULT_PROMPT_TRUNCATION_MODE, + # The production call from bco/ipo/kto/simpo caps only overflowing + # rows; this row overflows on the prompt, so the rebuild still runs. + only_when_overflow=True, + kl_chunk_size=kl_chunk_size, + )[0] + + for prefix in ("completion", "KL_completion"): + ids = list(capped[f"{prefix}_input_ids"]) + labels = list(capped[f"{prefix}_labels"]) + unmasked = [tok for tok, label in zip(ids, labels) if label != _IGNORE_LABEL] + assert unmasked == answer, (prefix, unmasked, answer) + assert unmasked.count(eos_id) == 1, (prefix, "EOS was doubled", unmasked) + assert len(ids) <= max_length, (prefix, len(ids)) + + +#: 41 words: longer than max_length // 2, yet prompt + completion is 45 of 64 tokens. +_FITTING_PROMPT = " ".join(["hi"] * 40 + ["world"]) +#: An over-long prompt whose end is visible, so the kept side of the cut can be asserted. +_OVERLONG_PROMPT_WITH_TAIL = " ".join(["hello"] * (_OVERLONG_PROMPT_WORDS - 1) + ["world"]) + + +def _setup_rows(tmp_path, monkeypatch, task, rows, val_rows=None): + import importlib + + module, cls_name, _ = _WRAPPERS[task] + weights = _tiny_llama_dir(tmp_path) + _write_tiny_tokenizer(weights) + monkeypatch.chdir(tmp_path) + cfg = _cfg(weights, tmp_path / "out", task) + wrapper = getattr(importlib.import_module(module), cls_name)(cfg, device="cpu") + dataset = {"train": rows} + if val_rows is not None: + dataset["val"] = val_rows + wrapper.setup(dataset) + return wrapper + + +def _rows_with(task, prompt, completion, n=8): + if task == "kto": + return [ + {"prompt": prompt, "completion": completion, "label": i % 2 == 0} + for i in range(n) + ] + return [{"prompt": prompt, "chosen": completion, "rejected": " bad"} for _ in range(n)] + + +class TestTheCapTouchesOnlyRowsThatOverflow: + @pytest.mark.parametrize("task", ("bco", "ipo", "kto", "simpo")) + def test_a_row_that_fits_reaches_the_trainer_uncut(self, tmp_path, monkeypatch, task): + rows = _rows_with(task, _FITTING_PROMPT, " good answer") + wrapper = _setup_rows(tmp_path, monkeypatch, task, rows) + row = wrapper.trainer.train_dataset[0] + prompt = list(row["prompt_ids"] if "prompt_ids" in row else row["prompt_input_ids"]) + words = wrapper.tokenizer.convert_ids_to_tokens(prompt) + assert len(prompt) > _MAX_PROMPT_LENGTH, (task, len(prompt), words[:2], words[-2:]) + assert words[-1] == "world", (task, words[-3:]) + + @pytest.mark.parametrize("task", ("bco", "ipo", "kto", "simpo")) + def test_a_validation_row_that_fits_reaches_the_trainer_uncut( + self, tmp_path, monkeypatch, task + ): + rows = _rows_with(task, "hi", " good answer") + val_rows = _rows_with(task, _FITTING_PROMPT, " good answer") + wrapper = _setup_rows(tmp_path, monkeypatch, task, rows, val_rows=val_rows) + row = wrapper.trainer.eval_dataset[0] + prompt = list(row["prompt_ids"] if "prompt_ids" in row else row["prompt_input_ids"]) + words = wrapper.tokenizer.convert_ids_to_tokens(prompt) + assert len(prompt) > _MAX_PROMPT_LENGTH, (task, len(prompt), words[:2], words[-2:]) + assert words[-1] == "world", (task, words[-3:]) + + @pytest.mark.parametrize("task", ("bco", "kto")) + def test_a_row_trl_already_fit_keeps_its_eos(self, tmp_path, monkeypatch, task): + rows = _rows_with(task, "hi", _LONG_COMPLETION) + wrapper = _setup_rows(tmp_path, monkeypatch, task, rows) + eos_id = wrapper.tokenizer.eos_token_id + for index, row in enumerate(wrapper.trainer.train_dataset): + if task == "bco" and len(row["answer_input_ids"]) < 10: + continue # BCO's rejected half (" bad") is not the long row + ids = list(row["completion_input_ids"]) + assert len(ids) <= _MAX_SEQUENCE_LENGTH, (task, index, len(ids)) + assert ids[-1] == eos_id, (task, index, "EOS dropped from a row trl already fit") + + @pytest.mark.parametrize("task", ("bco", "kto")) + def test_an_overlong_prompt_keeps_its_end_and_the_completion_its_eos( + self, tmp_path, monkeypatch, task + ): + rows = _rows_with(task, _OVERLONG_PROMPT_WITH_TAIL, " good answer") + wrapper = _setup_rows(tmp_path, monkeypatch, task, rows) + tok = wrapper.tokenizer + for index, row in enumerate(wrapper.trainer.train_dataset): + ids = list(row["completion_input_ids"]) + labels = list(row["completion_labels"]) + assert len(ids) <= _MAX_SEQUENCE_LENGTH, (task, index, len(ids)) + masked = [tok for tok, label in zip(ids, labels) if label == _IGNORE_LABEL] + assert tok.convert_ids_to_tokens(masked[-1:]) == ["world"], (task, index) + assert ids[-1] == tok.eos_token_id, (task, index, "the completion lost its EOS") + + def test_the_kto_kl_completion_keeps_its_answer(self, tmp_path, monkeypatch): + wrapper = _setup_rows(tmp_path, monkeypatch, "kto", _long_kto_rows(8)) + eos_id = wrapper.tokenizer.eos_token_id + for index, row in enumerate(wrapper.trainer.train_dataset): + ids = list(row["KL_completion_input_ids"]) + labels = list(row["KL_completion_labels"]) + assert len(ids) <= _MAX_SEQUENCE_LENGTH, (index, len(ids)) + non_eos = [ + token + for token, label in zip(ids, labels) + if label != _IGNORE_LABEL and token != eos_id + ] + assert non_eos, f"kto row {index}: the KL completion is EOS-only" + + def test_the_kto_kl_rebuild_matches_trl_on_rows_that_fit(self, tmp_path, monkeypatch): + """The KL reconstruction reproduces TRL's own rotated KL rows exactly. + + Rows that fit reach the trainer untouched, so they are TRL's own output; + rebuilding the KL completion from answer_input_ids with the same + per-device-batch rotation must give the identical ids, mask and labels. + """ + from soup_cli.trainer._trl_compat import _rebuild_unpaired_sequence + + completions = [" good answer", " fine", " a much longer answer with words", " no"] + rows = [ + {"prompt": "hi there", "completion": completions[index % 4], "label": index % 2 == 0} + for index in range(9) + ] + wrapper = _setup_rows(tmp_path, monkeypatch, "kto", rows) + dataset = wrapper.trainer.train_dataset + answers = list(dataset["answer_input_ids"]) + masks = list(dataset["answer_attention_mask"]) + chunk = wrapper.trainer.args.per_device_train_batch_size + eos_id = wrapper.tokenizer.eos_token_id + for index, row in enumerate(dataset): + lender = _trl_kl_lender(index, len(answers), chunk) + rebuilt = _rebuild_unpaired_sequence( + row["KL_prompt_input_ids"], + row["KL_prompt_attention_mask"], + answers[lender], + masks[lender], + eos_id, + max_length=_MAX_SEQUENCE_LENGTH, + max_prompt_length=_MAX_SEQUENCE_LENGTH, + truncation_mode="keep_end", + ) + assert rebuilt["input_ids"] == list(row["KL_completion_input_ids"]), index + assert rebuilt["labels"] == list(row["KL_completion_labels"]), index + assert rebuilt["attention_mask"] == list(row["KL_completion_attention_mask"]), index + + def test_the_kto_kl_completion_of_an_overlong_prompt_borrows_the_rotated_answer( + self, tmp_path, monkeypatch + ): + """An over-long row's KL completion is its prompt plus the neighbour's answer. + + TRL pairs each row with another row's answer inside every + per_device_train_batch_size chunk; the rebuild has to follow the same + rotation, or the KL estimate compares the wrong text. + """ + completions = [" alpha beta", " gamma", " delta epsilon zeta", " eta"] + rows = [ + { + "prompt": _OVERLONG_PROMPT_WITH_TAIL, + "completion": completions[index % 4], + "label": index % 2 == 0, + } + for index in range(8) + ] + wrapper = _setup_rows(tmp_path, monkeypatch, "kto", rows) + eos_id = wrapper.tokenizer.eos_token_id + dataset = wrapper.trainer.train_dataset + answers = list(dataset["answer_input_ids"]) + chunk = wrapper.trainer.args.per_device_train_batch_size + for index, row in enumerate(dataset): + lender = _trl_kl_lender(index, len(answers), chunk) + ids = list(row["KL_completion_input_ids"]) + labels = list(row["KL_completion_labels"]) + trained = [token for token, label in zip(ids, labels) if label != _IGNORE_LABEL] + assert len(ids) <= _MAX_SEQUENCE_LENGTH, (index, len(ids)) + assert trained == list(answers[lender]) + [eos_id], index + + +def _trl_kl_lender(index, size, chunk): + """trl's ``_get_kl_dataset`` rotates ``[a[-1]] + a[:-1]`` inside each chunk. + + So row ``i`` borrows row ``i - 1``'s answer, and the first row of a chunk + borrows the last one's. Written out here, not imported, so the test does + not share an oracle with the code under test. + """ + start = (index // chunk) * chunk + end = min(start + chunk, size) + return index - 1 if index > start else end - 1 + + +@pytest.mark.parametrize( + "prompt", ("hi there", _OVERLONG_PROMPT_WITH_TAIL), ids=("fits", "overlong") +) +def test_kto_kl_rows_borrow_the_predecessor_answer_at_batch_four(tmp_path, monkeypatch, prompt): + monkeypatch.setitem(_MIN_BATCH, "kto", _KL_ROTATION_BATCH) + # Distinct ids in the fixture vocabulary (unknown words all become ). + completions = [" good", " answer answer", " bad bad bad", " good bad good bad"] + rows = [ + {"prompt": prompt, "completion": completions[i % len(completions)], "label": i % 2 == 0} + for i in range(8) + ] + wrapper = _setup_rows(tmp_path, monkeypatch, "kto", rows) + assert wrapper.trainer.args.per_device_train_batch_size == _KL_ROTATION_BATCH + dataset = wrapper.trainer.train_dataset + answers = list(dataset["answer_input_ids"]) + eos_id = wrapper.tokenizer.eos_token_id + for index, row in enumerate(dataset): + ids = list(row["KL_completion_input_ids"]) + labels = list(row["KL_completion_labels"]) + trained = [tok for tok, lab in zip(ids, labels) if lab != _IGNORE_LABEL] + lender = _trl_kl_lender(index, len(answers), _KL_ROTATION_BATCH) + assert trained == list(answers[lender]) + [eos_id], index + + +@pytest.mark.parametrize("task", ("bco", "kto")) +def test_an_overlong_prompt_with_a_long_completion_fits_and_ends_in_eos( + tmp_path, monkeypatch, task +): + rows = _rows_with(task, _OVERLONG_PROMPT_WITH_TAIL, _LONG_COMPLETION) + wrapper = _setup_rows(tmp_path, monkeypatch, task, rows) + eos_id = wrapper.tokenizer.eos_token_id + prefixes = ["completion"] + (["KL_completion"] if task == "kto" else []) + for index, row in enumerate(wrapper.trainer.train_dataset): + for prefix in prefixes: + ids = list(row[f"{prefix}_input_ids"]) + labels = list(row[f"{prefix}_labels"]) + trained = [tok for tok, lab in zip(ids, labels) if lab != _IGNORE_LABEL] + assert len(ids) <= _MAX_SEQUENCE_LENGTH, (task, index, prefix, len(ids)) + assert ids[-1] == eos_id, (task, index, prefix, "EOS cut off by the cap") + assert len(trained) > 1, (task, index, prefix, "EOS-only completion") + + +def _world_id(wrapper): + return wrapper.tokenizer("world", add_special_tokens=False)["input_ids"][0] + + +@pytest.mark.parametrize("task", ("bco", "ipo", "kto", "simpo")) +def test_overlong_validation_rows_are_capped_too(tmp_path, monkeypatch, task): + rows = _rows_with(task, "hi", " good answer") + val_rows = _rows_with(task, _OVERLONG_PROMPT_WITH_TAIL, " good answer", n=4) + wrapper = _setup_rows(tmp_path, monkeypatch, task, rows, val_rows=val_rows) + eval_dataset = wrapper.trainer.eval_dataset + assert eval_dataset is not None and len(eval_dataset) > 0 + for row in eval_dataset: + if task == "ipo": + assert len(row["prompt_ids"]) + len(row["chosen_ids"]) <= _MAX_SEQUENCE_LENGTH + assert len(row["prompt_ids"]) + len(row["rejected_ids"]) <= _MAX_SEQUENCE_LENGTH + elif task == "simpo": + assert len(row["chosen_input_ids"]) <= _MAX_SEQUENCE_LENGTH + assert len(row["rejected_input_ids"]) <= _MAX_SEQUENCE_LENGTH + else: + assert len(row["completion_input_ids"]) <= _MAX_SEQUENCE_LENGTH + if task == "kto": + assert len(row["KL_completion_input_ids"]) <= _MAX_SEQUENCE_LENGTH + + +@pytest.mark.parametrize(("size", "chunk"), ((8, 4), (10, 4), (7, 3), (5, 2), (9, 2), (3, 8))) +def test_the_kl_rotation_stays_inside_its_chunk_and_handles_a_short_last_chunk(size, chunk): + from soup_cli.trainer._trl_compat import _rotated_answer_index + + for index in range(size): + assert _rotated_answer_index(index, size, chunk) == _trl_kl_lender(index, size, chunk) + + +def test_ipo_keeps_the_prompt_head_within_half_the_budget(tmp_path, monkeypatch): + rows = _rows_with("ipo", _OVERLONG_PROMPT_WITH_TAIL, " good answer") + wrapper = _setup_rows(tmp_path, monkeypatch, "ipo", rows) + world = _world_id(wrapper) + for row in wrapper.trainer.train_dataset: + assert len(row["prompt_ids"]) == _MAX_PROMPT_LENGTH + assert world not in list(row["prompt_ids"]) # keep_start drops the tail + + +def test_simpo_keeps_the_prompt_tail_within_half_the_budget(tmp_path, monkeypatch): + rows = _rows_with("simpo", _OVERLONG_PROMPT_WITH_TAIL, " good answer") + wrapper = _setup_rows(tmp_path, monkeypatch, "simpo", rows) + world = _world_id(wrapper) + for row in wrapper.trainer.train_dataset: + assert len(row["prompt_input_ids"]) == _MAX_PROMPT_LENGTH + assert list(row["prompt_input_ids"])[-1] == world # keep_end keeps the tail + + +def test_the_kto_prompt_masks_stay_as_long_as_their_ids(tmp_path, monkeypatch): + rows = _rows_with("kto", _OVERLONG_PROMPT_WITH_TAIL, " good answer") + wrapper = _setup_rows(tmp_path, monkeypatch, "kto", rows) + for row in wrapper.trainer.train_dataset: + assert len(row["prompt_attention_mask"]) == len(row["prompt_input_ids"]) + assert len(row["KL_prompt_attention_mask"]) == len(row["KL_prompt_input_ids"]) + + +def test_an_unpaired_row_that_exactly_fits_is_not_cut(): + from datasets import Dataset + + from soup_cli.trainer._trl_compat import enforce_preference_sequence_limit + + eos_id = 999 + prompt = list(range(100, 140)) # 40 tokens: longer than max_prompt_length 32 + answer = list(range(200, 223)) # 23 tokens: 40 + 23 + EOS == 64 == max_length + dataset = Dataset.from_dict( + { + "prompt_input_ids": [prompt], + "prompt_attention_mask": [[1] * len(prompt)], + "answer_input_ids": [answer], + "answer_attention_mask": [[1] * len(answer)], + "completion_input_ids": [prompt + answer + [eos_id]], + "completion_attention_mask": [[1] * 64], + "completion_labels": [[_IGNORE_LABEL] * len(prompt) + answer + [eos_id]], + } + ) + capped = enforce_preference_sequence_limit( + dataset, + max_length=64, + max_prompt_length=32, + truncation_mode="keep_end", + only_when_overflow=True, + )[0] + assert list(capped["completion_input_ids"]) == prompt + answer + [eos_id] + + +def test_a_dpo_layout_row_that_exactly_fits_is_not_cut(): + from datasets import Dataset + + from soup_cli.trainer._trl_compat import enforce_preference_sequence_limit + + prompt = list(range(100, 140)) # 40 tokens: longer than max_prompt_length 32 + chosen = list(range(200, 224)) # 40 + 24 == 64 == max_length + dataset = Dataset.from_dict( + {"prompt_ids": [prompt], "chosen_ids": [chosen], "rejected_ids": [list(range(300, 310))]} + ) + capped = enforce_preference_sequence_limit( + dataset, + max_length=64, + max_prompt_length=32, + truncation_mode="keep_start", + only_when_overflow=True, + )[0] + assert list(capped["prompt_ids"]) == prompt + assert list(capped["chosen_ids"]) == chosen + + +@pytest.mark.parametrize("task", ("bco", "kto")) +@pytest.mark.parametrize("prompt_kind", ("fits", "overlong")) +def test_an_answer_that_already_ends_in_eos_keeps_exactly_one( + tmp_path, monkeypatch, task, prompt_kind +): + prompt = "hi" if prompt_kind == "fits" else _OVERLONG_PROMPT_WITH_TAIL + rows = _rows_with(task, prompt, " good") + wrapper = _setup_rows(tmp_path, monkeypatch, task, rows) + eos_id = wrapper.tokenizer.eos_token_id + prefixes = ["completion"] + (["KL_completion"] if task == "kto" else []) + for row in wrapper.trainer.train_dataset: + for prefix in prefixes: + ids = list(row[f"{prefix}_input_ids"]) + labels = list(row[f"{prefix}_labels"]) + trained = [tok for tok, lab in zip(ids, labels) if lab != _IGNORE_LABEL] + assert trained[-1] == eos_id, (task, prompt_kind, prefix, trained) + assert trained.count(eos_id) == 1, (task, prompt_kind, prefix, trained) + + +@pytest.mark.parametrize("task", ("bco", "kto")) +def test_a_prompt_that_consumes_the_budget_still_keeps_a_short_answer( + tmp_path, monkeypatch, task +): + # The prompt alone is longer than ``max_length``, so TRL slices the short + # answer out of the completion. The conditional cap must still rebuild it + # (the sliced row overflows ``max_length``) rather than treat the row as + # fitting and skip it, which would train on an answer-less completion. + rows = _rows_with(task, _OVERLONG_PROMPT_WITH_TAIL, " good") + wrapper = _setup_rows(tmp_path, monkeypatch, task, rows) + eos_id = wrapper.tokenizer.eos_token_id + prefixes = ["completion"] + (["KL_completion"] if task == "kto" else []) + for row in wrapper.trainer.train_dataset: + for prefix in prefixes: + ids = list(row[f"{prefix}_input_ids"]) + labels = list(row[f"{prefix}_labels"]) + trained = [tok for tok, lab in zip(ids, labels) if lab != _IGNORE_LABEL] + assert len(ids) <= _MAX_SEQUENCE_LENGTH, (task, prefix, len(ids)) + assert trained[-1] == eos_id, (task, prefix, trained) + assert [t for t in trained if t != eos_id], (task, prefix, "answer discarded", trained)