Repository navigation
Fix preference max length enforcement across TRL wrappers - #1196
ousamabenyounes wants to merge 8 commits into
Conversation
MakazhanAlpamys
left a comment
MakazhanAlpamys
left a comment
There was a problem hiding this comment.
Review of 244b43fe.
Thanks for picking up #1185, and for covering all four uncapped wrappers instead of only the two I measured. I ran everything below on your head merged onto current main, CPU-only unless stated.
What holds, verified by running it
- Your RED/GREEN evidence reproduces exactly. On clean
mainyour new effective-batch test fails for bco, ipo, kto and simpo (4 failed, 2 passed): bcocompletion_input_ids(1, 163), ipoinput_ids(2, 164), ktocompletion_input_ids(2, 163), simpochosen_input_ids(1, 165). On your branch the contract file passes 21/21, and your five-file command passes 163. - Wider run: 842 passed, 14 skipped over 20 files that drive these wrappers. With the GPU visible,
test_v07204.py,test_issue1049_untied_streamed_preference.pyand the contract file pass 124/124, so the streamed simpo/kto path is fine. - Each call site is pinned. Disabling the new block in
simpo.py:227,kto.py:219,ipo.py:187andbco.py:239fails exactly one case, the one named for that wrapper. That meets the acceptance criterion in the issue. LeavingKL_prompt_*orKL_completion_*uncapped fails the kto case too. - No test was weakened. The ORPO helper extraction preserves behaviour: ORPO rows still carry
good answer </s>.
Blocking
-
The fragment fails the repository check, which will turn all nine test cells red.
tests/test_issue487_changelog_fragments.py::test_repository_fragments_match_the_newest_releasefails on your branch (1 failed, 20 passed; 21 passed onmain). The validator wants the text to start with-, and once that is fixed it also wants the text to reference the file's number. Fragments are also named after the PR. Please rungit mv changelog.d/0.75.1/1185.fixed.md changelog.d/0.75.1/1196.fixed.md, start the text with-, and end it with(#1185 by @ousamabenyounes in #1196). -
Every over-long KTO and BCO row trains on
</s>alone. The length is now bounded, but the completion is already gone. I decoded the prepared datasets from your own fixture (161-word prompt, completiongood answer) aftersetup(). In all 8 KTO rows and all 16 BCO rows, the only label that is not-100is EOS, and the KL completion looks the same. DPO, IPO, ORPO and SimPO keepgood answer </s>. The cause is upstream of your helper: trl slices the answer with a negative bound before the helper sees the row. That isanswer[:max_length - prompt_length]intrl/experimental/kto/kto_trainer.py:205-208, and the same slice intrl/experimental/bco/bco_trainer.py:288-290. Both files are byte-identical in trl 0.29.0 and 0.29.1, so this covers the whole supported range. Onmainthose rows are 163 tokens and have no completion. This PR fixes the length but not the missing completion. The prompt-longer-than-max_lengthcase is the acceptance scenario in #1185, and the project's rule is to refuse or report, never to drop content silently (RAFT drops such rows with a warning,docs/training.md:1523). A fix is available. The prepared dataset still holds the untruncated answer inanswer_input_ids(I checked:['good', 'answer'],['bad'], and all 80 ids for an 80-token answer) and the untruncated prompt inprompt_input_ids. Socap_unpairedcan rebuildcompletion_*from those two columns (capped prompt, then answer up to the remaining budget, then EOS the way trl appends it, with-100over the prompt) instead of re-slicing trl'scompletion_*.KL_*has no answer column to rebuild from, so either rebuild it the same way or raise an error that names the row count andmax_length. -
The test cannot see the content of a row, and that is why item 2 passes. It asserts only
shape[-1] <= 64and the row count. Each of these mutations still passes all 21 tests:completion_limit = 0at_trl_compat.py:171, which deletes every completion;completion_limit = 10**9, which never caps a completion;- leaving the prompt tokens trainable in
*_labelsat:178; - flipping the prompt/completion boundary test at
:162; - passing
max_prompt_length = max_lengthinkto.py.
The
10**9mutation survives because every fixture completion is one or two words, so the completion cap on the new unpaired path never runs. Please assert per prepared row that the tokens whose label is not-100equal the expected truncated completion and include at least one non-EOS token, and add a long-completion row for each layout. One trap when you do:key.endswith("input_ids")(test file line 256) also matchesanswer_input_ids. That is an unused trl column that stays untruncated. With a long completion, a correct KTO row hascompletion_input_ids(2, 64) butanswer_input_ids(2, 80), so the assertion as written would fail it. Limit the check to the model inputs. -
KTO truncates the prompt at the wrong end.
kto.py:223passesDEFAULT_PROMPT_TRUNCATION_MODE = "keep_start", which drops the end of the prompt. When I checked, the prompt's last word was gone. With a chat template, the end of the prompt is the generation header (<|im_start|>assistant\n), so the completion gets trained after a cut-off user turn. trl 0.29'sKTOConfighas notruncation_modefield. Its default before the removal waskeep_end(trl 0.26.2experimental/kto/kto_config.py:195-196), andBCOConfig,CPOConfigandORPOConfigstill default tokeep_endin 0.29.1. As it stands, KTO and BCO share a layout but cut opposite ends. Please usekeep_endfor KTO and add a test that pins the direction. Flipping it today leaves all 21 tests passing.
A decision I owed you, made: cap only when the row does not fit. As written, the cap is unconditional, so any prompt longer than max_length // 2 gets cut even when the whole row fits. On a 45-of-64-token row, main keeps the 42-token prompt and this branch cuts it to 32, for all four wrappers. You followed the issue, which said to mirror DPO and ORPO, and those two are unconditional on main; that wording was mine. But trl itself cut KTO, BCO and CPO prompts only when prompt plus response overran max_length, and cutting a row that fits throws away training data for nothing. So for the four wrappers here, please apply the cap only when the row exceeds max_length, and add one case with a row that fits and must reach the trainer unchanged. Whether DPO and ORPO should become conditional too is a separate question, and I will open an issue for it rather than widen this PR.
Non-blocking
- A related problem is already on
main, and this PR doesn't touch it; I'll file it separately. Whenchosenis longer thanmax_length, trl's CPO sliceanswer[: max_length - longer_response_length](cpo_trainer.py:513-515) leaves a shortrejectedwith no tokens. SimPO's length-normalised log-prob then becomes 0/0. After two steps all 8 LoRA tensors were NaN on bothmainand this branch, and the logged loss read 0.0, because transformers filters NaN out of the log by default (with the filter off, the losses read[nan, nan]). Both control fixtures trained normally. If your fix for item 2 moves truncation ahead of trl's tokenizer, this may come along for free. It doesn't have to. - IPO inherits DPO's
keep_startprompt cap. That matches DPO onmain, so any direction change belongs in a DPO follow-up, not here. - Docs: no page says how the preference tasks truncate prompts. One line in the preference section of
docs/training.mdwould cover all six tasks: the cap ismax_length // 2, name which end is kept, and the completion gets the rest.
Nits
DEFAULT_PROMPT_TRUNCATION_MODEreads like a module-wide default, but only KTO uses it. Name it for KTO, or inline it once item 4 is done.- The blank line after
if orpo_columns <= columns:was dropped; the other branches keep theirs. - A one-line docstring on
_cap_combined_preference_sequencewould help: it treats the first label that is not-100as the prompt/completion boundary, and the boundary mutation above shows that line matters.
|
@ousamabenyounes, following up on my 09-24 review of
The one open design question is already answered on my side: cap only when a row overflows Through 2026-10-01 it stays yours, and a short "still on it" keeps it with nothing else owed. If there's no word by then, I'll close this without prejudice and mark #1185 as available again, so the next person can see it is free. That's not a judgement on the work already here — the RED/GREEN evidence and covering all four wrappers instead of two were both right. |
Addresses the review on MakazhanAlpamys#1196: - cap_unpaired rebuilds completion_* from the untruncated answer_input_ids / prompt_input_ids. TRL's KTO/BCO tokenizer slices the answer to fit max_length before Soup's cap runs, so an over-length prompt left completion_input_ids with only the trailing EOS; re-deriving from TRL's own truncated column could never recover the answer. - DEFAULT_PROMPT_TRUNCATION_MODE switches to keep_end. KTOConfig has no truncation_mode field, so KTO caps its prompt through this default; keep_end matches TRL's historical KTO/BCO/CPO/ORPO direction and preserves a chat template's generation header at the prompt tail. - Tests: per-row assertion that the unmasked completion tokens equal the truncated answer (including a non-EOS token), a long-completion fixture, and a keep_end direction pin. The content check reads completion_labels, never the unused answer_input_ids column. - Rename the changelog fragment to 1196.fixed.md and give it the required list item and attribution.
|
Thanks — all four are addressed in b1ca232.
Conditional capping kept as you settled it — only when a row overflows |
MakazhanAlpamys
left a comment
MakazhanAlpamys
left a comment
There was a problem hiding this comment.
@ousamabenyounes, thanks for turning this around so quickly. Review of b1ca232e. I ran everything below on it merged onto current main (d7a4f7a7), CPU only, with trl 0.29.1.
What this push fixed, verified by running it
- The KTO/BCO completion is back. On an over-long prompt (160 ×
hello+world, completiongood answer),mainhas 8 of 8 KTO rows and 16 of 16 BCO rows whose only trainable token is</s>. On your branch that is 0 of 8 and 0 of 16, and the trainable tokens decode togood answer </s>. - The new content test has real teeth. The five changes that passed all 21 tests last time now fail it: deleting every completion, never capping the completion, leaving the prompt trainable, inverting the prompt/completion boundary, and
max_prompt_length = max_lengthinkto.py. It also catches disabling your rebuild (your RED), unmasking the prompt inside the rebuilt row, and dropping the answer. Across the 19 changes I made to the code, 16 fail your contract file. - KTO keeps the end of the prompt now, and setting
DEFAULT_PROMPT_TRUNCATION_MODEback tokeep_startfails your new unit test. - The fragment passes the validator. The contract file plus the fragment tests pass 45 of 45. A wider run passes 756 with 14 skipped; it covers 23 related files, including the six preference trainers, the dispatcher, the streaming suite and the recent GRPO tests. ruff is clean. On
main, your contract file fails 7 of 24, so every new test is red without the fix. - DPO and ORPO come out identical to
mainon all three fixtures I used, which is right while #1209 is still open.
Still blocking
-
The cap is still unconditional. Your reply says conditional capping is in place, but a row that fits is still cut. The fixture: a 41-word prompt ending in
worldplus a 2-word completion, 45 of 64 tokens.mainhands KTO, BCO and SimPO a 42-token prompt and IPO a 41-token one; this branch cuts all four to 32. IPO losesworld, and the end of the prompt is where a chat template puts the assistant header. KTO, BCO and SimPO lose the leading<s>. The condition in_truncate_tokens(_trl_compat.py:155) asks whether the prompt is longer thanmax_length // 2. The review asked whether the row is longer thanmax_length, plus one case with a row that fits.One constraint on the fix: IPO shares DPO's
prompt_idslayout, and all six wrappers call the same helper. Please make "only when the row overflows" an argument that the four call sites pass (simpo.py:227,kto.py:219,ipo.py:187,bco.py:239). Leavedpo.pyandorpo.pyexactly as they are, so that this PR does not decide #1209 as a side effect. -
KTO/BCO rows that trl had already fit now lose their EOS (new in this push). Take a short prompt with a 100-token completion. trl truncates the answer and then appends EOS (
trl/experimental/kto/kto_trainer.py:202-208, then:232-238), so onmainthe row is 64 tokens ending in</s>. On this branch,_rebuild_unpaired_completion_rowrebuilds the full prompt, the full answer and the EOS, and the cap (_trl_compat.py:183-186) then slices the EOS off. The row is 64 tokens ending inanswer, and that happens to every long-completion KTO and BCO row. That contradicts the helper's own docstring ("never drops the stop signal"). KTO's KL rows still end in</s>, so the matched and KL completions now disagree. Rebuilding only when the original row overflows fixes this together with item 1. When you do rebuild, truncate the answer to the budget minus one and then append the EOS, as trl does. -
KTO's KL rows are still EOS-only when the prompt overflows. This is the second half of item 2 in the 09-24 review ("
KL_*has no answer column to rebuild from, so either rebuild it the same way or raise an error that names the row count andmax_length"). My 09-27 summary left it out, so that part is on me.- With
_long_kto_rows, 8 of 8KL_completionrows train on</s>only: 33 tokens, one trainable. _mixed_kto_rowshides this. Each long prompt's KL row borrows its neighbour's 100-token completion, and trl's negative slice leaves 1 of those 100 tokens, so a non-EOS check passes there.- Fix path: trl builds the KL rows by rotating
answer_input_idsby one inside eachper_device_train_batch_sizechunk (kto_trainer.py:85-93,:597-603), andanswer_input_idssurvives untruncated. So you can rebuild the KL completion the same way you rebuilt the matched one. Pin it with a test showing that your reconstruction equals trl's own KL rows exactly on a fixture that fits. - Or raise an error that names the row count and
max_length.
- With
Tests that pin all three. These are ready to paste at the end of tests/test_trl_preference_config_contract.py; they use only that file's own helpers.
#: 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):
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")
wrapper.setup({"train": rows})
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", "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 = [t for t, lab in zip(ids, labels) if lab == _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 = [t for t, lab in zip(ids, labels) if lab != _IGNORE_LABEL and t != eos_id]
assert non_eos, f"kto row {index}: the KL completion is EOS-only"I ran this exact block appended to your file, and it is ruff-clean.
- On this branch, 7 of the 9 new cases fail (4 fits, 2 EOS, 1 KL) and 26 tests pass.
- On
main, the 4 fits and 2 EOS cases pass, because they pin whatmainalready does; the rest fail. - A fix for items 1-3 should pass all 33.
The third test also catches two changes your file currently lets through. One is kto.py:223 passing "keep_start" literally instead of the constant. The other is the rebuild never appending the EOS. Each of those leaves all 24 contract tests green, and 228 tests across six files.
Non-blocking
keep_endnow drops the leading<s>from every capped KTO row. BCO and SimPO on this branch, and ORPO onmain, do the same. trl adds BOS after its own truncation (kto_trainer.py:218-230), so its cut never removed it. Keeping a leading BOS and taking the tail of the rest would match trl. Since ORPO behaves this way already, this can be a follow-up.- Docs, still open from 09-24: one line in the preference section of
docs/training.mdsaying when the prompt is cut and which end is kept. - Nits:
- The blank line after
if orpo_columns <= columns:(_trl_compat.py:279) is still missing. _cap_combined_preference_sequence(:162) still has no docstring.- The fragment has no final period, while most fragments in
changelog.d/0.75.1/end with one. Once item 1 lands, it can also say that prompts are cut only when a row exceedsmax_length.
- The blank line after
Your update keeps #1185 with you, as the 09-27 note said. CI on b1ca232e has all 13 required checks queued, and none has started yet.
MakazhanAlpamys
left a comment
MakazhanAlpamys
left a comment
There was a problem hiding this comment.
@ousamabenyounes, a status check on b1ca232e, and a first answer to a question the red board will have raised: the CI failures on this PR are not yours.
- Every one of the 16 jobs on this head is
cancelled, none isfailed. Each annotation reads "The run was canceled by @MakazhanAlpamys", at 2026-09-28 07:56 UTC, and no job ever got a runner (no steps, no log). That was a clean-up of queued runs whose head already carried a changes-requested review, to drain a CI backlog; it says nothing about your code. - The only real CI failure this PR ever had was on the previous head: the fragment check (
fragment must start with a Markdown list item). You fixed it inb1ca232e. - A re-run of the cancelled jobs would reuse a merge ref frozen on 09-27, well over a hundred commits behind
main, so the next push is what will give you a meaningful run.
I also re-ran your head myself, merged with current main (89dc498a, clean merge), CPU only, trl 0.29.1:
- Your contract file plus the fragment tests: 45 passed. 43 test files that import the six wrappers,
_trl_compat, the dispatch, layer streaming and the MoE/seed/kwarg-drift scans: 1764 passed, 26 skipped, 0 failed.ruffis clean. - The claim in #1185 reproduces on unmodified
main: withmax_length: 64and your 161-token prompt, the real batch is[1, 165]for SimPO,[2, 163]for KTO,[2, 164]for IPO and[1, 163]for BCO. Your contract file fails 7 of 24 onmainand passes 24 of 24 here. - Seven mutations of your fix each fail your file: the rebuild branch off (2 failed), any one of the four call sites off (1, 2, 1 and 2),
keep_startfor KTO (1), andcompletion_limit = 0(2). The tests do pin what they claim to pin.
The three items from my 09-28 review are still open, and I checked each by running it. I appended the nine cases from that review to your file on this head: 7 fail, 26 pass, the same as on the 28th.
- The cap is unconditional: a row that fits (41-word prompt, 45 of 64 tokens) is still cut to a 32-token prompt for all four wrappers.
- KTO and BCO rows trl had already fit lose their trailing
</s>in the rebuild. - KTO's KL completions are still
</s>only when the prompt overflows.
I went one step further than last time so that the next round is short. I wrote a patch for all three and ran it: only_when_overflow passed from the four call sites (dpo.py and orpo.py untouched, so #1209 stays open), a rebuild that cuts the answer to max_length - len(prompt) - 1 and appends the EOS after the cut, and a KL rebuild that follows trl's rotation of answer_input_ids inside each per_device_train_batch_size chunk. Results on this head merged with main (d2b0d153; the diff also applies cleanly on 89dc498a):
- your contract file with the nine cases from that review and two more of mine (KL rows equal trl's own rotated rows exactly on rows that fit, and an over-long row's KL completion is its prompt plus the neighbour's answer): 35 passed;
- the same 43 wrapper files: 1775 passed, 26 skipped, 0 failed;
ruffclean; - six mutations of the patch, each caught: the gate off (2 failed), the EOS not appended (4), the rotation replaced by the identity (1), the SimPO call site (1),
kl_chunk_sizedropped inkto.py(1),keep_start(2).
The diff is attached below. It is a proposal, not something to paste blindly: git apply takes it on your head, and you are welcome to restructure it. I have not run the GPU or streamed paths, only CPU.
Still small, for the same push: one line in the preference section of docs/training.md saying when the prompt is cut and which end is kept; the fragment can say prompts are cut only when a row exceeds max_length, and end with a period; and the two nits from last time (the blank line after if orpo_columns <= columns:, a docstring on _cap_combined_preference_sequence).
This stays yours through 2026-10-01, as noted on the 27th; a short "still on it" keeps it with nothing else owed.
Patch (applies to this head merged with current main)
diff --git a/src/soup_cli/trainer/_trl_compat.py b/src/soup_cli/trainer/_trl_compat.py
index cddeb9819..d0460132d 100644
--- a/src/soup_cli/trainer/_trl_compat.py
+++ b/src/soup_cli/trainer/_trl_compat.py
@@ -191,44 +191,61 @@ def _cap_combined_preference_sequence(
}
-def _rebuild_unpaired_completion_row(row: dict[str, Any]) -> 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.
+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]]:
+ """One KTO/BCO sequence rebuilt from the untruncated prompt and answer.
+
+ TRL slices the answer with a negative bound once the prompt alone reaches
+ ``max_length`` (``answer[:max_length - len(prompt)]``), so the ``*completion_*``
+ columns of an over-long row hold an EOS and nothing else. The prompt and the
+ answer survive untruncated in ``prompt_input_ids`` / ``answer_input_ids``, so
+ cut the prompt to ``max_prompt_length``, give the answer what is left of
+ ``max_length`` minus one, and append the EOS AFTER that cut, as TRL does.
"""
- prompt_ids = list(row["prompt_input_ids"])
- prompt_mask = list(row["prompt_attention_mask"])
- answer_ids = list(row["answer_input_ids"])
- answer_mask = list(row["answer_attention_mask"])
- eos_id = list(row["completion_input_ids"])[-1]
+ prompt_ids = _truncate_tokens(list(prompt_ids), max_prompt_length, truncation_mode)
+ prompt_mask = _truncate_tokens(list(prompt_mask), max_prompt_length, truncation_mode)
+ answer_ids = list(answer_ids)
+ answer_mask = list(answer_mask)
if answer_ids and answer_ids[-1] == eos_id:
- eos_ids: list[int] = []
- eos_mask: list[int] = []
- else:
- eos_ids = [eos_id]
- eos_mask = [1] # EOS is a real, attended token
+ answer_ids, answer_mask = answer_ids[:-1], answer_mask[:-1]
+ budget = max(0, max_length - len(prompt_ids) - 1)
+ body, body_mask = answer_ids[:budget], answer_mask[:budget]
return {
- "completion_input_ids": prompt_ids + answer_ids + eos_ids,
- "completion_attention_mask": prompt_mask + answer_mask + eos_mask,
- "completion_labels": [_IGNORE_INDEX] * len(prompt_ids) + answer_ids + eos_ids,
+ "input_ids": prompt_ids + body + [eos_id],
+ "attention_mask": prompt_mask + body_mask + [1],
+ "labels": [_IGNORE_INDEX] * len(prompt_ids) + body + [eos_id],
}
+def _rotated_answer_index(index: int, size: int, chunk: int) -> int:
+ """Row whose answer TRL's ``_get_kl_dataset`` lends to row ``index``.
+
+ TRL rotates ``answer_input_ids`` by one inside each ``per_device_train_batch_size``
+ chunk of the dataset (the last chunk may be shorter), so a row borrows its
+ predecessor's answer and the first row of a chunk borrows the last one's.
+ """
+ start = (index // chunk) * chunk
+ end = min(start + chunk, size)
+ return index - 1 if index > start else end - 1
+
+
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.
@@ -240,6 +257,13 @@ def enforce_preference_sequence_limit(
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`` leaves every row whose TRL-built sequences already
+ fit in ``max_length`` untouched (TRL itself cut KTO/BCO/CPO prompts only
+ when a row overran), so a prompt longer than ``max_prompt_length`` is cut
+ only when the row does not fit. DPO and ORPO keep the unconditional cap.
+ ``kl_chunk_size`` is the per-device train batch size TRL rotated KTO's KL
+ answers within; it is needed to rebuild an over-long KL row.
"""
columns = set(getattr(dataset, "column_names", ()))
dpo_columns = {"prompt_ids", "chosen_ids", "rejected_ids"}
@@ -264,6 +288,10 @@ def enforce_preference_sequence_limit(
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
)
@@ -277,7 +305,13 @@ def enforce_preference_sequence_limit(
return dataset.map(cap_dpo)
if orpo_columns <= columns:
+
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(
list(row["prompt_input_ids"]), max_prompt_length, truncation_mode
)
@@ -306,61 +340,94 @@ def enforce_preference_sequence_limit(
return dataset.map(cap_orpo)
if unpaired_columns <= columns:
-
- def cap_unpaired(row: dict[str, Any]) -> dict[str, Any]:
+ has_answer = {"answer_input_ids", "answer_attention_mask"} <= columns
+ has_kl = {
+ "KL_prompt_input_ids",
+ "KL_prompt_attention_mask",
+ "KL_completion_input_ids",
+ "KL_completion_attention_mask",
+ "KL_completion_labels",
+ } <= columns
+ answers = list(dataset["answer_input_ids"]) if has_answer and has_kl else []
+ answer_masks = list(dataset["answer_attention_mask"]) if answers else []
+ 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 {}
+ eos_id = list(row["completion_input_ids"])[-1]
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
)
- # 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)
- if {"answer_input_ids", "answer_attention_mask"} <= columns
- else row
- )
- capped = {
+ capped: dict[str, Any] = {
"prompt_input_ids": prompt_ids,
"prompt_attention_mask": prompt_mask,
- **_cap_combined_preference_sequence(
- completion_source,
- "completion",
+ }
+ if has_answer:
+ completion = _rebuild_unpaired_sequence(
+ row["prompt_input_ids"],
+ row["prompt_attention_mask"],
+ row["answer_input_ids"],
+ row["answer_attention_mask"],
+ eos_id,
max_length=max_length,
max_prompt_length=max_prompt_length,
truncation_mode=truncation_mode,
- ),
- }
- 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:
+ capped["completion_input_ids"] = completion["input_ids"]
+ capped["completion_attention_mask"] = completion["attention_mask"]
+ capped["completion_labels"] = completion["labels"]
+ else:
capped.update(
_cap_combined_preference_sequence(
row,
- "KL_completion",
+ "completion",
max_length=max_length,
max_prompt_length=max_prompt_length,
truncation_mode=truncation_mode,
)
)
+ if has_kl:
+ 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 answers:
+ lender = _rotated_answer_index(index, len(answers), chunk)
+ kl = _rebuild_unpaired_sequence(
+ row["KL_prompt_input_ids"],
+ row["KL_prompt_attention_mask"],
+ answers[lender],
+ 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)
+ 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 4e9b913f8..33d0d796e 100644
--- a/src/soup_cli/trainer/bco.py
+++ b/src/soup_cli/trainer/bco.py
@@ -243,6 +243,7 @@ class BCOTrainerWrapper:
"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
diff --git a/src/soup_cli/trainer/ipo.py b/src/soup_cli/trainer/ipo.py
index 665436c68..b3f7f9614 100644
--- a/src/soup_cli/trainer/ipo.py
+++ b/src/soup_cli/trainer/ipo.py
@@ -191,6 +191,7 @@ class IPOTrainerWrapper:
"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
diff --git a/src/soup_cli/trainer/kto.py b/src/soup_cli/trainer/kto.py
index 9e1e69dcf..2b9868102 100644
--- a/src/soup_cli/trainer/kto.py
+++ b/src/soup_cli/trainer/kto.py
@@ -226,6 +226,8 @@ class KTOTrainerWrapper(StreamingSetupMixin):
"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
diff --git a/src/soup_cli/trainer/simpo.py b/src/soup_cli/trainer/simpo.py
index e42a6b8ae..2469c3982 100644
--- a/src/soup_cli/trainer/simpo.py
+++ b/src/soup_cli/trainer/simpo.py
@@ -231,6 +231,7 @@ class SimPOTrainerWrapper(StreamingSetupMixin):
"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_trl_preference_config_contract.py b/tests/test_trl_preference_config_contract.py
index 8af120acd..fdc5d363b 100644
--- a/tests/test_trl_preference_config_contract.py
+++ b/tests/test_trl_preference_config_contract.py
@@ -481,3 +481,151 @@ class TestKtoTruncatesThePromptFromTheEnd:
if label == _IGNORE_LABEL
]
assert masked == tail
+
+
+#: 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):
+ 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")
+ wrapper.setup({"train": rows})
+ 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", "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 = [t for t, lab in zip(ids, labels) if lab == _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 = [t for t, lab in zip(ids, labels) if lab != _IGNORE_LABEL and t != 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,
+ _rotated_answer_index,
+ )
+
+ completions = [" good answer", " fine", " a much longer answer with words", " no"]
+ rows = [
+ {"prompt": "hi there", "completion": completions[i % 4], "label": i % 2 == 0}
+ for i 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 = _rotated_answer_index(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.
+ """
+ from soup_cli.trainer._trl_compat import _rotated_answer_index
+
+ completions = [" alpha beta", " gamma", " delta epsilon zeta", " eta"]
+ rows = [
+ {
+ "prompt": _OVERLONG_PROMPT_WITH_TAIL,
+ "completion": completions[i % 4],
+ "label": i % 2 == 0,
+ }
+ for i 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 = _rotated_answer_index(index, len(answers), chunk)
+ 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]
+ assert len(ids) <= _MAX_SEQUENCE_LENGTH, (index, len(ids))
+ assert trained == list(answers[lender]) + [eos_id], index
|
@ousamabenyounes, thank you for the quick rounds on #1196: the KTO/BCO completion is back and the content tests pin it, which I verified by running them. Three items and a few small ones are still open on
Because your checks were cancelled on our side during a queue clean-up, the window I set earlier for this PR is extended: it stays yours through 2026-10-04. |
|
@ousamabenyounes, a short check-in on #1196. What is already solid: rebuilding the KTO/BCO completion from Still open (details in my review of 28 September; the review of 30 September has a tested patch for 1-3 that
The nine ready-to-paste cases in the 28 September review pin items 1-3. A fix for all three should pass the whole contract file. The CI jobs on this head were cancelled during our queue clean-up; your code did not cause it. Re-running them would only re-test the old merge from 27 September. Your next push starts a fresh run. The branch still merges cleanly with current If you are still on it, a one-line note is enough. If you would rather hand it on, just say so, and a maintainer or another contributor can pick it up with your work credited. |
|
Pushed What changed:
Validation:
|
|
Current head is I reran the focused review/changelog command in a clean CPU validation venv: 36 passed. I also reran |
|
Pushed and verified current head What changed since the last review round:
Validation:
|
There was a problem hiding this comment.
Thanks, @ousamabenyounes. The three previous requests are addressed on 4c9e8d0: predecessor rotation, IPO's documented keep_start default, and the EOS budget regression. All 44 amended contract cases pass in a native integration with main 3b28d69.
[P2] Preserve an answer's existing terminal EOS during rebuilding.
File: src/soup_cli/trainer/_trl_compat.py:213-215
An overlong prompt plus the valid short completion good</s> reaches _rebuild_unpaired_sequence with raw answer IDs [good, EOS]. The helper retains those IDs and appends another EOS, so BCO's matched target, KTO's matched target and KTO's KL target train on [good, EOS, EOS]. The original TRL path appends EOS only when absent; a fitting prompt preserves the single EOS. This changes the target for inputs that already include their stop token.
Please preserve an already-terminal EOS while enforcing the answer budget, for both matched and KL reconstruction, and add caller-level regressions that compare exact unmasked tokens against a known vocabulary oracle. Keep the existing fitting-row, no-explicit-EOS and long-answer controls; no broader tokenizer change is needed.
Independent verification used real BCO/KTO wrapper setup and the existing tiny local tokenizer/checkpoint: the new external file gives exactly 3 failures and 9 passing controls. A temporary narrow production-source control that removes the existing terminal EOS before the current cap/append makes the unchanged 12 cases pass. That control was removed, every byte restored, full Ruff passed, and the native merge aborted cleanly. The authored 44-case file remains unchanged.
Evidence is CPU-only on Torch 2.14.0, Transformers 5.17.0 and TRL 0.29.1. Available physical RAM exceeded 6 GiB before each child. Each child verified its owned source and ran offline with one CPU thread, disabled CUDA visibility and a 90-second bound. No model download, provider, GPU test or training step was performed. This review does not claim a fresh full matrix or repeat the historical rotation/budget mutations.
|
Heads-up for the merge, @ousamabenyounes: v0.76.0 is out, so changelog fragments now go under |
4a3edc9 to
8874d7c
Compare
|
Done, thanks for the heads-up — pushed
I did not replay the whole suite locally (37,853 collected tests is several hours on my 4-core box), so remote CI on |
Addresses the review on MakazhanAlpamys#1196: - cap_unpaired rebuilds completion_* from the untruncated answer_input_ids / prompt_input_ids. TRL's KTO/BCO tokenizer slices the answer to fit max_length before Soup's cap runs, so an over-length prompt left completion_input_ids with only the trailing EOS; re-deriving from TRL's own truncated column could never recover the answer. - DEFAULT_PROMPT_TRUNCATION_MODE switches to keep_end. KTOConfig has no truncation_mode field, so KTO caps its prompt through this default; keep_end matches TRL's historical KTO/BCO/CPO/ORPO direction and preserves a chat template's generation header at the prompt tail. - Tests: per-row assertion that the unmasked completion tokens equal the truncated answer (including a non-EOS token), a long-completion fixture, and a keep_end direction pin. The content check reads completion_labels, never the unused answer_input_ids column. - Rename the changelog fragment to 1196.fixed.md and give it the required list item and attribution.
|
@ousamabenyounes, thanks for the rebase and the fragment move. I checked
Both tests are in my review of 5 October, ready to paste at the end of the contract file. With these three in and the required checks green, the maintainers will queue it. |
|
@MakazhanAlpamys, thanks for the precise review. All three items are addressed. Current head:
Validation (CPU, trl 0.29.1):
|
|
@ousamabenyounes, thank you for the fast round. The three items from my 13 October comment are closed, and I checked each one by running it rather than by reading your reply:
What I ran (CPU only, head 1. An answer that already ends in EOS gets a second one (this is the capped_prompt_mask = _truncate_tokens(list(prompt_mask), max_prompt_length, truncation_mode)
answer_ids, answer_mask = list(answer_ids), list(answer_mask)
if answer_ids and answer_ids[-1] == eos_id:
# TRL appends EOS only when it is absent: keep exactly one
answer_ids, answer_mask = answer_ids[:-1], answer_mask[:-1]
answer_limit = max(0, max_length - len(capped_prompt_ids) - 1)2. Fourteen of the 52 mutations still pass the file, 12 of them real gaps. Each line below passes all 44 tests when broken:
The two others ( This block, appended to 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</s>")
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)After both of these I found nothing else open: the merge with |
4c9e8d0 to
8208aa0
Compare
Addresses the review on MakazhanAlpamys#1196: - cap_unpaired rebuilds completion_* from the untruncated answer_input_ids / prompt_input_ids. TRL's KTO/BCO tokenizer slices the answer to fit max_length before Soup's cap runs, so an over-length prompt left completion_input_ids with only the trailing EOS; re-deriving from TRL's own truncated column could never recover the answer. - DEFAULT_PROMPT_TRUNCATION_MODE switches to keep_end. KTOConfig has no truncation_mode field, so KTO caps its prompt through this default; keep_end matches TRL's historical KTO/BCO/CPO/ORPO direction and preserves a chat template's generation header at the prompt tail. - Tests: per-row assertion that the unmasked completion tokens equal the truncated answer (including a non-EOS token), a long-completion fixture, and a keep_end direction pin. The content check reads completion_labels, never the unused answer_input_ids column. - Rename the changelog fragment to 1196.fixed.md and give it the required list item and attribution.
|
The red
Both look like test updates on this branch rather than code changes, but that is yours and the maintainer's call. For what it is worth, #1209 is stacked on |
|
@ousamabenyounes, this is a review of The 1. Two red tests. They are the only two failures in
On the rule itself: for SimPO, KTO, IPO and BCO a prompt is cut only when the row does not fit in 2. Text that now says the opposite. Please change:
3. The test block from 19 October. None of the cases in that block is in A realistic path. Everything above is a small edit or a paste. A suggested order: paste the 19 October block, rewrite the SimPO boundary test, drop line 114, fix the wording in the four places, then run Timing, kindly. Nothing here is urgent and nothing you wrote is lost. Your last push was on 26 October and this list is new, so we count from today. Our queue rule is that a pull request waiting on its author for 14 days gets a ping in the thread, and if another 7 days pass we close it with a note that it can be reopened at any time. Counted from today that is 9 November and 16 November. A short "still on it" is enough to tell us you are here; if your time is short, say so and a maintainer or another contributor can finish the last steps with your work credited. Two other pieces wait behind this pull request: #1209 (the same rule for DPO and ORPO) and the CUDA run for #840. |
|
still on it - thank for your review will send you update soon |
|
Thank you for the word, @ousamabenyounes, take the time you need. One small update to the review above, because v0.77.0 was released today: the |
Addresses the review on MakazhanAlpamys#1196: - cap_unpaired rebuilds completion_* from the untruncated answer_input_ids / prompt_input_ids. TRL's KTO/BCO tokenizer slices the answer to fit max_length before Soup's cap runs, so an over-length prompt left completion_input_ids with only the trailing EOS; re-deriving from TRL's own truncated column could never recover the answer. - DEFAULT_PROMPT_TRUNCATION_MODE switches to keep_end. KTOConfig has no truncation_mode field, so KTO caps its prompt through this default; keep_end matches TRL's historical KTO/BCO/CPO/ORPO direction and preserves a chat template's generation header at the prompt tail. - Tests: per-row assertion that the unmasked completion tokens equal the truncated answer (including a non-EOS token), a long-completion fixture, and a keep_end direction pin. The content check reads completion_labels, never the unused answer_input_ids column. - Rename the changelog fragment to 1196.fixed.md and give it the required list item and attribution.
v0.76.0 is released, so the fragment belongs under changelog.d/0.76.0/. Keeping it under 0.75.1 fails test_repository_fragments_match_the_newest_release once this branch contains main.
TRL builds KL rows with [a[-1]] + a[:-1] inside each batch chunk, so row i borrows row i-1's answer; _rotated_answer_index returned i+1. Both directions coincide at batch 2, which hid the bug from the existing tests. - follow TRL's predecessor rotation in _rotated_answer_index - KL tests use an independent oracle and a batch-4 case that separates the two directions - pin the EOS slot in answer_limit with an over-long prompt plus long completion case - docs: IPO follows DPO's truncation_mode (keep_start by default) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…aired rows _rebuild_unpaired_sequence appended EOS unconditionally, so a valid short completion whose answer_input_ids already end in the stop token ([good, </s>]) was rebuilt as [good, </s>, </s>] on the matched and KL targets whenever an overlong prompt forced the rebuild. TRL appends EOS only when absent; split the terminal EOS off the answer body before enforcing the budget and re-appending. Adds a per-row content regression that compares the unmasked completion and KL tokens against the exact answer oracle and asserts a single EOS.
The cap now runs only when a prepared row overflows data.max_length for SimPO/KTO/IPO/BCO, so the two boundary tests that reached main (MakazhanAlpamys#2116 SimPO prompt-budget, MakazhanAlpamys#2111 KTO source check) encoded the old unconditional rule and turned the test job red. Rewrite the SimPO boundary case to assert an over- budget prompt is left whole when the row fits and cut to max_length // 2 only when it overflows, drop the stale KTO source assertion, and add the review's coverage block pinning the eval-split cap, KL rotation chunk bounds, prompt budget/direction, mask lengths, exact-fit no-op, and single-EOS rebuild. Move the changelog fragment to 0.77.0 and sync the SimPO/KTO prose in docs/training.md and docs/backends-and-ops.md to the conditional rule.
8208aa0 to
cac3cc9
Compare
|
Rebased onto main; fixed conditional trimming, answer preservation and EOS/KL handling, added the requested coverage tests, and updated the docs and 0.77.0 changelog. RED/GREEN verified; 308 tests passed. |
Summary
data.max_lengthenforcement for SimPO and KTO when TRL no longer acceptsmax_prompt_lengthanswer_input_idswhen TRL already sliced the completion awaymax_length; DPO/ORPO stay unchanged for Preference prompt cap: should DPO and ORPO cut the prompt only when the row does not fit? #1209Closes #1185.
Test verification (RED -> GREEN)
8874d7c's production code:python -m pytest tests/test_trl_preference_config_contract.py -o addopts="" -q-> 1 failed (test_kto_kl_rows_borrow_the_predecessor_answer_at_batch_four[overlong]: KL rows borrowed the successor's answer, trl uses the predecessor), 43 passed.4c9e8d02527466c62cf2c14842cc6e71d49a4ae5: same command -> 44 passed; dropping the- 1inanswer_limitfailstest_an_overlong_prompt_with_a_long_completion_fits_and_ends_in_eos[bco|kto].4c9e8d0: ruff clean on the CI paths, 15 related test files 519 passed, Python diff coverage 100% (87/87 changed production lines); 569 passed on a local merge with currentmain.4c9e8d0:lint,plan,gate, andtest (ubuntu-latest, 3.12)succeeded; the remaining jobs were skipped by workflow conditions.