Skip to content

Fix negative truncation bound in experimental BCO _process_tokens - #6581

Closed
ErenAta16 wants to merge 1 commit into
huggingface:mainfrom
ErenAta16:fix-bco-truncation-negative-bound
Closed

ErenAta16 wants to merge 1 commit into
huggingface:mainfrom
ErenAta16:fix-bco-truncation-negative-bound

Conversation

@ErenAta16

@ErenAta16 ErenAta16 commented Jul 29, 2026 •

Copy link
Copy Markdown

What does this PR do?

_process_tokens in the experimental BCO trainer truncates the response with
answer[: max_length - len(prompt_input_ids)]. Nothing caps the prompt in BCOTrainer, so that
bound goes negative as soon as the prompt is longer than the (BOS/EOS adjusted) max_length.
answer[:negative] slices from the end of the response, which empties it outright whenever the
response is shorter than the overflow, and the sequence stays over max_length anyway because the
prompt was never shortened.

Calling _process_tokens directly with the tiny test tokenizer:

 prompt  answer  max_len |  seq_len  <=max?  answer_tokens
     10      10       32 |       22    True             11
     28      10       32 |       32    True              3
     60      10       32 |       62   False              1
     60      40       32 |       72   False             11
   1000      10       32 |     1002   False              1

answer_tokens: 1 is the appended EOS and nothing else, so those rows carry no completion tokens
at all and contribute no learning signal, while the 1000-token one still costs a full forward pass
under a max_length of 32.

Where it came from. #4964 removed max_prompt_length from BCO. Before it the branch was two
steps: truncate the prompt to max_prompt_length, then give the response whatever was left. The
response bound got the better - len(prompt_input_ids) form, but the step that kept the prompt
bounded went away with max_prompt_length and nothing replaced it. That PR's description says
"only the response is truncated if the total sequence exceeds max_length", which cannot hold when
the prompt on its own is already past max_length.

The change. Restore the two-step order, with max_length as the prompt cap instead of the
removed max_prompt_length. answer_budget can no longer be negative and the sequence honours
max_length in every case. This is the policy KTOTrainer applies when it truncates the
concatenated sequence with full_ids[: self.max_length].

Same table after the change:

 prompt  answer  max_len |  seq_len  <=max?  answer_tokens
     10      10       32 |       22    True             11
     28      10       32 |       32    True              3
     60      10       32 |       32    True              1
     60      40       32 |       32    True              1
   1000      10       32 |       32    True              1

Rows 3 to 5 still end up with no completion tokens, because there is no way to honour max_length
and keep completion tokens when the prompt alone fills the budget. That matches KTO. What changes
is that the sequence is now within budget and the response is never cut from the end.

Question for maintainers. DPOTrainer goes further and drops the rows where truncation leaves
nothing behind (dpo_trainer.py, the .filter() after tokenization). Adding the same filter to BCO
would be a two-line change next to the existing desirable/undesirable filters, but it changes the row
count and therefore the desirable-to-undesirable ratio BCO relies on, so it felt like your call. Happy
to add it here or separately.

Tests. test_process_tokens_truncation_respects_max_length in
tests/experimental/test_bco_trainer.py, parametrized over five prompt/answer length combinations.
It asserts the sequence stays within max_length, that both halves remain prefixes of what went in
(a negative bound keeps a suffix instead), and that the response survives intact when it fits.

Reverting only bco_trainer.py fails the [60-10] and [1000-10] cases:

FAILED tests/experimental/test_bco_trainer.py::TestBCOTrainer::test_process_tokens_truncation_respects_max_length[60-10]
FAILED tests/experimental/test_bco_trainer.py::TestBCOTrainer::test_process_tokens_truncation_respects_max_length[1000-10]
2 failed, 3 passed

The three that pass either way are guards against over-truncating the cases that already worked.

$ pytest tests/experimental/test_bco_trainer.py -k "tokenize_and_process_tokens or truncation_respects_max_length"
6 passed

CPO and ORPO have a different flavour of the same problem in tokenize_row (they size the bound off
longer_response_length instead of the prompt), tracked in #6548. That one is being picked up
separately; this PR does not touch either file.

Fixes # (issue)

Before submitting

  • This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
  • Did you read the contributor guideline, Pull Request section?
  • Was this discussed/approved via a GitHub issue? Please add a link to it if that's the case.
  • Did you make sure to update the documentation with your changes?
  • Did you write any new necessary tests?

AI writing disclosure

  • No AI usage: the PR was written entirely by a human.
  • AI-assisted: some parts were suggested or improved by AI, but the PR was written and reviewed by a human.
  • AI-generated: the PR was mostly or fully generated by an AI tool.

Who can review?

@qgallouedec @albertvillanova — flagging you since #4964 is the change this follows up on.


Note

Medium Risk
Changes tokenization for long prompts in experimental BCO training, which can alter which examples contribute completion tokens but fixes incorrect sequences and wasted forward passes.

Overview
Fixes overlong sequences in the experimental BCO trainer’s _process_tokens when the prompt alone exceeds max_length.

Previously only the answer was capped with answer[: max_length - len(prompt)]. If the prompt was longer than the budget, that slice bound went negative, Python kept a suffix of the answer (often wiping real completion tokens) while the prompt was never shortened, so completion_input_ids could still exceed max_length.

The change truncates the prompt prefix first (to max_length), then applies the remaining budget to the answer—aligned with KTOTrainer’s concatenated truncation. Comments document the old failure mode.

Adds test_process_tokens_truncation_respects_max_length, parametrized over five prompt/answer lengths, checking total length ≤ max_length, prefix-only truncation for prompt and answer, and intact answers when everything fits.

Reviewed by Cursor Bugbot for commit 81c9e86. Bugbot is set up for automated code reviews on this repo. Configure here.

`_process_tokens` truncates the response with `answer[: max_length - len(prompt_input_ids)]`.
Nothing caps the prompt, so that bound goes negative as soon as the prompt is
longer than the (BOS/EOS adjusted) `max_length`. A negative bound slices from the
*end* of the response instead of the start, which empties the response outright
whenever it is shorter than the overflow, and the sequence stays over `max_length`
regardless because the prompt was never shortened:

    prompt=  10 answer= 10 max_length=  32 -> completion_len=  22   answer tokens: 11
    prompt=  28 answer= 10 max_length=  32 -> completion_len=  32   answer tokens:  3
    prompt=  60 answer= 10 max_length=  32 -> completion_len=  62   answer tokens:  1  (EOS only)
    prompt=1000 answer= 10 max_length=  32 -> completion_len=1002   answer tokens:  1  (EOS only)

The last two carry no real completion tokens, so the example contributes nothing
but still costs a full forward pass at 1000+ tokens with `max_length=32`.

This is what huggingface#4964 left behind. Before that PR the branch truncated the prompt to
`max_prompt_length` first and only then gave the response what was left; removing
`max_prompt_length` dropped the first step, and the replacement bound has nothing
bounding the prompt any more.

Restore the two-step order with `max_length` as the prompt cap, so the response
budget can no longer be negative and the sequence honours `max_length`. This is
what `KTOTrainer` does when it truncates the concatenation with
`full_ids[: max_length]`.

An example whose prompt fills `max_length` still ends up with no completion tokens,
same as KTO. `DPOTrainer` additionally drops those rows from the dataset; doing the
same here would change the desirable/undesirable balance BCO relies on, so I left it
out of this change.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants