[behroozazarkhalili]

Resolves #6996

BCOTrainer fails before the first optimizer step on any encoder-decoder model. Two independent
defects sit on that path, and the first masks the second, so one run only ever shows one of them.

Defect 1: the collator has no case for the answer_input_ids column

Measured on main at d9e880c:

File "trl/experimental/utils.py", line 94, in __call__
    raise ValueError(f"Unexpected key in batch '{k}'")
ValueError: Unexpected key in batch 'answer_input_ids'

_tokenize (bco_trainer.py:228-233) adds answer_input_ids and answer_attention_mask to every
example, and embedding_input_ids too when an embedding tokenizer is given. The encoder-decoder
branch of DPODataCollatorWithPadding (utils.py:81-94) assigned a padding value to
prompt*_input_ids, to any *_attention_mask, and to keys starting with chosen, rejected or
completion or containing decoder. Both new columns matched none of those and reached the final
raise.

The decoder-only branch of the same collator (utils.py:96-113) already pads every *_input_ids
column with the pad token, which is why the failure was specific to encoder-decoder models. This
matches on the suffix there as well rather than listing prefixes, so the two branches stay in step
and the next column someone adds is covered. Reordering the checks leaves every key that already
worked on its old padding value.

Defect 2: forward() reads a key the encoder-decoder path never builds

Patch past defect 1 and the same script reaches:

File "trl/experimental/bco/bco_trainer.py", line 1189, in forward
    batch["completion_input_ids"],
KeyError: 'completion_input_ids'

_process_tokens builds completion_input_ids only on its decoder-only branch
(bco_trainer.py:304). Its encoder-decoder branch (bco_trainer.py:344-352) builds
prompt_input_ids, prompt_attention_mask, completion_labels, completion_attention_mask and
completion_decoder_input_ids, and no concatenated input ids, because the encoder reads the prompt
and the completion arrives as labels. forward() now reads the prompt on that branch, which is what
the sibling compute_reference_log_probs (bco_trainer.py:1088-1114) already does.

Tests

tests/experimental/test_bco_trainer.py had no encoder-decoder case at all. The one test that
passed an is_encoder_decoder kwarg read it off a decoder-only trainer, so the value was always
False.

Two tests now cover the path. test_train_encoder_decoder runs it plainly.
test_train_encoder_decoder_udm runs it with an embedding function, which reaches the collator from
_get_sample_prompt_embeddings during init rather than from the training loop, and is the only test
that catches the embedding_* column.

Both training tests only show that a run completes: a collator that padded the answer_* and
embedding_* ids with -100 passes them, because those ids feed the reward embedding and the UDM
classifier, where a wrong pad token changes a number rather than raising.
test_collator_pads_answer_and_embedding_ids_with_pad_token calls the collator on a ragged
encoder-decoder batch and pins the values: those ids take the pad token, their masks take 0, and
completion_labels keeps -100.

Verification

Check	Result
Both tests, fix present	2 passed
Revert the collator fix	both fail, ValueError: Unexpected key in batch 'answer_input_ids'
Revert the forward() fix	both fail, KeyError: 'completion_input_ids'
Narrow the collator rule to a ("prompt", "answer") allowlist	plain test passes, UDM test fails on embedding_input_ids
Collator test, fix present	1 passed
Collator test against main's collator	fails, ValueError: Unexpected key in batch 'answer_input_ids'
Pad answer_* and embedding_* with -100	both training tests pass, collator test fails
ruff 0.13.3 check and format --check	clean
doc-builder style, pinned rev, max_len 119	0 changes

The third mutant is the one that shows the UDM test is not redundant: it is the only case that
distinguishes a rule covering answer_* from one covering every token column.

Scope

Neither defect depends on the checkpoint. Both sites are gated only on self.is_encoder_decoder, so
both fire for every encoder-decoder model. I measured one,
trl-internal-testing/tiny-T5ForConditionalGeneration.

DPODataCollatorWithPadding is shared with CPO, ORPO, OnlineDPO, NashMD and XPO. Those trainers keep
their answer_* values as locals rather than dataset columns, so none of them reached the raise, and
their suites were run to confirm the reorder leaves them alone.

The *_decoder_input_ids keys share the -100 rule, and a model would reject -100 as a token id.
They are only built when a model is passed to _process_tokens (and to CPO's and ORPO's
tokenize_row), which no trainer does, so on every trainer path the key is absent and the model
derives its decoder inputs from the labels. That is pre-existing and unchanged here; the comment on
the rule now says so.

---

Note

Medium Risk
Touches shared DPODataCollatorWithPadding encoder-decoder logic and BCO training forward; behavior is gated on is_encoder_decoder and other DPO-family trainers were checked for regressions.

Overview
BCOTrainer could not train encoder-decoder models (e.g. T5): batching failed on answer_input_ids / embedding_input_ids, and forward() then looked for completion_input_ids, which the encoder-decoder tokenization path never creates.

The encoder-decoder branch of DPODataCollatorWithPadding now assigns padding like the decoder-only path: masks use 0, completion/label fields use -100, and remaining *_input_ids (prompt, answer_*, embedding_*) use the tokenizer pad id—so BCO batches no longer hit Unexpected key.

BCOTrainer.forward on encoder-decoder runs the model on prompt_input_ids with completion_labels (and optional decoder_input_ids), matching compute_reference_log_probs instead of concatenated completion ids.

New tests cover a full T5 BCO train step, UDM with embeddings at init-time collator use, and explicit collator padding assertions for answer_* / embedding_*.

^{Reviewed by Cursor Bugbot for commit f2eb458. Bugbot is set up for automated code reviews on this repo. Configure here.}

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[qgallouedec]

Closing in favour of #7136, which removes trl.experimental.bco. Both defects you reported are real, there is just no code left for the fix to land on. thks for the digging.