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

tests/test_kto_trainer.py

@@ -203,7 +203,7 @@ def test_padding_and_masks(self):
         torch.testing.assert_close(result["KL_input_ids"], expected_kl_input_ids)
         torch.testing.assert_close(result["KL_attention_mask"], expected_kl_attention_mask)
         torch.testing.assert_close(result["KL_completion_mask"], expected_kl_completion_mask)
-        assert result["label"] == [True, False]
+        torch.testing.assert_close(result["label"], torch.tensor([True, False]))
 
     def test_optional_reference_logps(self):
         collator = DataCollatorForUnpairedPreference(pad_token_id=0)
@@ -932,6 +932,24 @@ def test_train_with_iterable_dataset(self):
             new_param = trainer.model.get_parameter(n)
             assert not torch.equal(param, new_param), f"Parameter {n} has not changed."
 
+    def test_evaluate_with_iterable_dataset(self):
+        # Regression test: Accelerate's dispatch for `IterableDataset` requires every batch field to be a tensor, so
+        # the collator must not leave `"label"` as a plain list of booleans, which broke streaming evaluation with
+        # `TypeError: Can only concatenate tensors but got <class 'bool'>`.
+        train_dataset = load_dataset("trl-internal-testing/zen", "standard_unpaired_preference", split="train")
+        eval_dataset = load_dataset(
+            "trl-internal-testing/zen", "standard_unpaired_preference", split="test", streaming=True
+        )
+
+        # `apo_zero_unpaired` avoids KTO's KL term, which is not possible with a streaming dataset.
+        training_args = KTOConfig(output_dir=self.tmp_dir, loss_type="apo_zero_unpaired", report_to="none")
+        trainer = KTOTrainer(
+            model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", args=training_args, train_dataset=train_dataset
+        )
+
+        metrics = trainer.evaluate(eval_dataset=eval_dataset)
+        assert metrics["eval_loss"] is not None
+
     def test_train_with_chat_template_kwargs(self):
         dataset = load_dataset("trl-internal-testing/zen", "conversational_preference", split="train")
 

trl/trainer/kto_trainer.py

@@ -102,7 +102,7 @@ class DataCollatorForUnpairedPreference(DataCollatorMixin):
     When `"KL_completion_ids"` is present, the same three tensors are returned for the (mismatched) KL sequence under
     the `"KL_input_ids"`, `"KL_attention_mask"` and `"KL_completion_mask"` keys.
 
-    The returned dictionary also contains a `"label"` key with the list of labels. Optionally, the examples can contain
+    The returned dictionary also contains a `"label"` key with a tensor of labels. Optionally, the examples can contain
     `"ref_logps"` and `"ref_KL_logps"` keys, in which case the returned dictionary will also contain these keys with
     the corresponding tensors.
 
@@ -133,7 +133,7 @@ class DataCollatorForUnpairedPreference(DataCollatorMixin):
                                [1, 1, 1, 0, 0]]),
      'completion_mask': tensor([[0, 0, 0, 1, 1],
                                 [0, 0, 1, 0, 0]]),
-     'label': [True, False]}
+     'label': tensor([ True, False])}
 
     >>> # With KL completions
     >>> examples = [
@@ -153,7 +153,7 @@ class DataCollatorForUnpairedPreference(DataCollatorMixin):
                                   [1, 1, 1, 1]]),
      'KL_completion_mask': tensor([[0, 0, 0, 1],
                                    [0, 0, 1, 1]]),
-     'label': [True, False]}
+     'label': tensor([ True, False])}
     ```
     """
 
@@ -205,7 +205,8 @@ def torch_call(self, examples: list[dict[str, Any]]) -> dict[str, Any]:
             output["ref_logps"] = torch.tensor([example["ref_logps"] for example in examples])
         if "ref_KL_logps" in examples[0]:
             output["ref_KL_logps"] = torch.tensor([example["ref_KL_logps"] for example in examples])
-        output["label"] = [example["label"] for example in examples]
+        # Must be a tensor: Accelerate cannot concatenate non-tensor fields across processes (accelerate#4111)
+        output["label"] = torch.tensor([example["label"] for example in examples])
         return output
 
 
@@ -231,7 +232,7 @@ class DataCollatorForVisionUnpairedPreference(DataCollatorMixin):
     - `"attention_mask"`: Tensor indicating attention mask.
     - `"pixel_values"`: Tensor representing image pixel values.
     - `"completion_mask"`: Tensor indicating which tokens correspond to completions.
-    - `"label"`: List of booleans indicating whether each completion is desirable.
+    - `"label"`: Tensor of booleans indicating whether each completion is desirable.
     - When `calculate_kl` is `True`: `"KL_input_ids"`, `"KL_attention_mask"` and `"KL_completion_mask"` for the cycled
       KL sequences.
 
@@ -294,7 +295,7 @@ class DataCollatorForVisionUnpairedPreference(DataCollatorMixin):
                                   [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]),
      'KL_completion_mask': tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0],
                                    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1]]),
-     'label': [True, False]}
+     'label': tensor([ True, False])}
     ```
     """
 
@@ -454,7 +455,8 @@ def torch_call(self, examples: list[dict[str, Any]]) -> dict[str, Any]:
             if "mm_token_type_ids" in processed_prompts:
                 output["KL_mm_token_type_ids"] = kl_mm_token_type_ids
 
-        output["label"] = [example["label"] for example in examples]
+        # Must be a tensor: Accelerate cannot concatenate non-tensor fields across processes (accelerate#4111)
+        output["label"] = torch.tensor([example["label"] for example in examples])
         return output