Fix KTO evaluate crash on streaming datasets with unpaired-preference data

[albertvillanova]

This PR fixes streaming evaluation for KTOTrainer with unpaired-preference data, where a batch field returned as a plain list broke Accelerate's dispatch mechanism for IterableDataset.

See: https://github.com/huggingface/trl/actions/runs/28934157283/job/85840257000?pr=6326

FAILED tests/test_kto_trainer.py::TestKTOTrainer::test_evaluate_with_eval_dataset_dict[True] - TypeError: Can only concatenate tensors but got <class 'bool'>

Traceback:

  >       metrics = trainer.evaluate(eval_dataset=eval_dataset)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  
  tests/test_kto_trainer.py:369: 
  _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
  trl/trainer/kto_trainer.py:1648: in evaluate
      return super().evaluate(
  .venv/lib/python3.12/site-packages/transformers/trainer.py:4451: in evaluate
      dataset_metrics = self.evaluate(
  trl/trainer/kto_trainer.py:1648: in evaluate
      return super().evaluate(
  .venv/lib/python3.12/site-packages/transformers/trainer.py:4469: in evaluate
      output = eval_loop(
  .venv/lib/python3.12/site-packages/transformers/trainer.py:4655: in evaluation_loop
      for step, inputs in enumerate(dataloader):
                          ^^^^^^^^^^^^^^^^^^^^^
  .venv/lib/python3.12/site-packages/accelerate/data_loader.py:856: in __iter__
      next_batch, next_batch_info = self._fetch_batches(main_iterator)
                                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  .venv/lib/python3.12/site-packages/accelerate/data_loader.py:812: in _fetch_batches
      batch = concatenate(batches, dim=0)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  .venv/lib/python3.12/site-packages/accelerate/utils/operations.py:617: in concatenate
      return type(data[0])({k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  .venv/lib/python3.12/site-packages/accelerate/utils/operations.py:615: in concatenate
      return honor_type(data[0], (concatenate([d[i] for d in data], dim=dim) for i in range(len(data[0]))))
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  .venv/lib/python3.12/site-packages/accelerate/utils/operations.py:81: in honor_type
      return type(obj)(generator)
             ^^^^^^^^^^^^^^^^^^^^
  .venv/lib/python3.12/site-packages/accelerate/utils/operations.py:615: in <genexpr>
      return honor_type(data[0], (concatenate([d[i] for d in data], dim=dim) for i in range(len(data[0]))))
                                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
  
  data = [False], dim = 0
  
      def concatenate(data, dim=0):
          """
          Recursively concatenate the tensors in a nested list/tuple/dictionary of lists of tensors with the same shape.
      
          Args:
              data (nested list/tuple/dictionary of lists of tensors `torch.Tensor`):
                  The data to concatenate.
              dim (`int`, *optional*, defaults to 0):
                  The dimension on which to concatenate.
      
          Returns:
              The same data structure as `data` with all the tensors concatenated.
          """
          if isinstance(data[0], (tuple, list)):
              return honor_type(data[0], (concatenate([d[i] for d in data], dim=dim) for i in range(len(data[0]))))
          elif isinstance(data[0], Mapping):
              return type(data[0])({k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
          elif not isinstance(data[0], torch.Tensor):
  >           raise TypeError(f"Can only concatenate tensors but got {type(data[0])}")
  E           TypeError: Can only concatenate tensors but got <class 'bool'>
  
  .venv/lib/python3.12/site-packages/accelerate/utils/operations.py:619: TypeError

Motivation

DataCollatorForUnpairedPreference and DataCollatorForVisionUnpairedPreference returned the "label" batch field as a plain Python list of booleans instead of a tensor, while every other field in the batch was a tensor. This works fine with a map-style Dataset, but Accelerate's DataLoaderDispatcher, used for streaming IterableDatasets, needs every batch field to be a tensor so it can concatenate and redistribute sub-batches across processes. Passing a streaming eval dataset to KTOTrainer.evaluate() therefore raised TypeError: Can only concatenate tensors but got <class 'bool'>.

Note this is not an Accelerate-version-specific issue. Accelerate ≥1.13.0 added a short-circuit in concatenate() that happens to skip the crash when num_processes == 1 (see huggingface/accelerate#4111, which the Accelerate maintainers confirm is a partial, single-process-only fix). Real multi-process/distributed streaming evaluation still crashes on every Accelerate version, old and new. So this fix is required unconditionally, regardless of installed Accelerate version.

Solution

"label" is now built with torch.tensor(...) in both collators, consistent with every other returned field. Docstrings and doctest examples were updated to reflect the new tensor output.

Changes

• Return "label" as a tensor in DataCollatorForUnpairedPreference.torch_call and DataCollatorForVisionUnpairedPreference.torch_call
• Update docstrings and doctest examples for both collators to show "label" as a tensor
• Update the existing test_padding_and_masks assertion for the new tensor type
• Add a regression test that runs a real streaming evaluation through KTOTrainer.evaluate() with an IterableDataset
• Add a code comment on both collator changes referencing concatenate() still raises TypeError on non-tensor batch fields when num_processes > 1 accelerate#4111, explaining why this must stay unconditional

---

Note

Low Risk
Narrow collator output change; loss paths already accept tensor labels via torch.tensor/torch.as_tensor. Risk is mainly if any custom code assumed label was always a list.

Overview
Fixes KTOTrainer.evaluate crashing on streaming (IterableDataset) eval when Accelerate’s dataloader dispatch tries to merge batch fields that must all be tensors.

DataCollatorForUnpairedPreference and DataCollatorForVisionUnpairedPreference now set label with torch.tensor([...]) instead of a Python list of booleans, matching other batch keys and satisfying Accelerate’s concatenation (see accelerate#4111). Docstrings and doctest examples were updated accordingly.

Tests assert tensor labels in the collator padding test and add test_evaluate_with_iterable_dataset, which runs a real streaming eval with apo_zero_unpaired (no KL term on streaming data).

^{Reviewed by Cursor Bugbot for commit 7929cf2. 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.

[albertvillanova]

Why this isn't gated on an Accelerate version

While debugging the tests_min_versions CI failure for this fix, I initially assumed it was Accelerate-version-specific, since the regular tests job (latest Accelerate) passes while tests_min_versions (accelerate==1.4.0) fails. Wanted to record the investigation here in case anyone else looks at this and wonders the same thing.

Turns out that's misleading: Accelerate 1.13.0 added a short-circuit in concatenate() that returns a single gathered batch as-is instead of type-checking it.

• Allow non-Tensor values in a batch with dispatch_batches=True accelerate#3850

DataLoaderDispatcher._fetch_batches gathers exactly num_processes batches before calling concatenate, so this short-circuit only fires when num_processes == 1, i.e. a plain, non-distributed test run, which is all our CI exercises. It masks the crash for single-process runs on Accelerate ≥1.13.0, but does nothing for num_processes > 1.

Confirmed with the Accelerate maintainers themselves in the PR review thread (huggingface/accelerate#3850 (comment)): they explicitly acknowledged multi-process support was out of scope for that fix

"it'll always fail in multi-process settings... multi-process support is not possible".

I opened a GH issue to track this as a documented, discoverable limitation:

• concatenate() still raises TypeError on non-tensor batch fields when num_processes > 1 accelerate#4111

So: real distributed streaming evaluation with KTO's unpaired-preference collator would still crash today, on any Accelerate version, without this fix. The fix here ("label" as a tensor, matching every other collated field) is the correct, version-independent solution, not something to conditionally apply only for old Accelerate.