Fix mixed-image Online DPO generation

tests/experimental/test_online_dpo_trainer.py

@@ -12,12 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from types import SimpleNamespace
+
 import pytest
+import torch
 from datasets import Dataset, DatasetDict, features, load_dataset
 from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer
 from transformers.utils import is_peft_available, is_vision_available
 
-from trl.experimental.online_dpo import OnlineDPOConfig, OnlineDPOTrainer
+from trl.experimental.online_dpo import OnlineDPOConfig, OnlineDPOTrainer, online_dpo_trainer
 
 from ..testing_utils import TrlTestCase, require_peft, require_torch_accelerator, require_vision, require_vllm
 
@@ -32,6 +35,150 @@
     from transformers import AutoModelForImageTextToText, AutoProcessor
 
 
+def test_training_step_prepares_mixed_image_batch_for_vllm_server(monkeypatch):
+    class StopAfterGeneration(Exception):
+        pass
+
+    class DummyProcessor:
+        def __call__(self, **kwargs):
+            self.text = kwargs["text"]
+            return {"input_ids": torch.tensor([[1], [2], [3]])}
+
+    class DummyClient:
+        def chat(self, messages, **kwargs):
+            self.messages = messages
+            return {"completion_ids": [[4], [5], [6], [7], [8], [9]]}
+
+    monkeypatch.setattr(
+        online_dpo_trainer,
+        "apply_chat_template",
+        lambda example, processor: {
+            "prompt": " ".join(part["type"] for message in example["prompt"] for part in message["content"])
+        },
+    )
+    monkeypatch.setattr(online_dpo_trainer, "gather_object", lambda value: value)
+    monkeypatch.setattr(online_dpo_trainer, "broadcast_object_list", lambda value, from_process: value)
+
+    trainer = OnlineDPOTrainer.__new__(OnlineDPOTrainer)
+    trainer.state = SimpleNamespace(global_step=0)
+    trainer._move_model_to_vllm = lambda: None
+    trainer._tokenizer = SimpleNamespace(eos_token_id=0, pad_token_id=0)
+    trainer.processing_class = DummyProcessor()
+    trainer.accelerator = SimpleNamespace(is_main_process=True, process_index=0, device="cpu")
+    trainer.num_generations = 2
+    trainer.repetition_penalty = 1.0
+    trainer.temperature = 1.0
+    trainer.top_p = 1.0
+    trainer.top_k = None
+    trainer.min_p = None
+    trainer.generation_config = SimpleNamespace(max_tokens=16)
+    trainer.args = SimpleNamespace(generation_kwargs=None, use_vllm=True)
+    trainer.vllm_client = DummyClient()
+    trainer.vllm_mode = "server"
+
+    image = object()
+    embedded_image = object()
+    prompts = [
+        [{"role": "user", "content": "Describe it"}],
+        [{"role": "user", "content": "Say hello"}],
+        [
+            {
+                "role": "user",
+                "content": [{"type": "image", "image": embedded_image}, {"type": "text", "text": "Embedded"}],
+            }
+        ],
+    ]
+    generated_prompts = None
+    generate = trainer._generate_vllm
+
+    def generate_then_stop(prompts, images):
+        nonlocal generated_prompts
+        generated_prompts = prompts
+        generate(prompts, images)
+        raise StopAfterGeneration
+
+    trainer._generate_vllm = generate_then_stop
+
+    with pytest.raises(StopAfterGeneration):
+        trainer.training_step(SimpleNamespace(train=lambda: None), {"prompt": prompts, "image": [image, None, None]})
+
+    assert generated_prompts[0][0]["content"][0] == {"type": "image", "image": image}
+    assert generated_prompts[1][0]["content"] == [{"type": "text", "text": "Say hello"}]
+    assert generated_prompts[2][0]["content"][0] == {"type": "image", "image": embedded_image}
+    assert prompts[0][0]["content"] == "Describe it"
+    assert prompts[1][0]["content"] == "Say hello"
+
+    assert trainer.vllm_client.messages[0][0]["content"][0] == {"type": "image", "image": image}
+    assert trainer.vllm_client.messages[1][0]["content"] == [{"type": "text", "text": "Say hello"}]
+    assert trainer.vllm_client.messages[2][0]["content"][0] == {"type": "image", "image": embedded_image}
+    assert trainer.processing_class.text == ["image text", "text", "image text"]
+
+
+def test_generate_transformers_with_mixed_image_batch(monkeypatch):
+    class StopAtProcessor(Exception):
+        pass
+
+    class DummyProcessor:
+        def __call__(self, **kwargs):
+            self.images = kwargs["images"]
+            raise StopAtProcessor
+
+    monkeypatch.setattr(
+        online_dpo_trainer, "maybe_apply_chat_template", lambda example, processor: {"prompt": "prompt"}
+    )
+
+    trainer = OnlineDPOTrainer.__new__(OnlineDPOTrainer)
+    trainer._tokenizer = SimpleNamespace(eos_token_id=0, pad_token_id=0)
+    trainer.processing_class = DummyProcessor()
+    trainer.image_token = None
+    image = object()
+
+    with pytest.raises(StopAtProcessor):
+        trainer._generate(torch.nn.Linear(1, 1), ["image prompt", "text prompt"], [[image], []])
+
+    assert trainer.processing_class.images == [[image], []]
+
+
+def test_generate_vllm_colocate_with_mixed_image_batch(monkeypatch):
+    class DummyLLM:
+        def generate(self, inputs, generation_config, use_tqdm):
+            self.inputs = inputs
+            return [
+                SimpleNamespace(
+                    outputs=[SimpleNamespace(token_ids=[3]), SimpleNamespace(token_ids=[4])], prompt_token_ids=[1]
+                ),
+                SimpleNamespace(
+                    outputs=[SimpleNamespace(token_ids=[5]), SimpleNamespace(token_ids=[6])], prompt_token_ids=[2]
+                ),
+            ]
+
+    monkeypatch.setattr(
+        online_dpo_trainer,
+        "apply_chat_template",
+        lambda example, processor: {
+            "prompt": " ".join(part["type"] for message in example["prompt"] for part in message["content"])
+        },
+    )
+
+    trainer = OnlineDPOTrainer.__new__(OnlineDPOTrainer)
+    trainer.args = SimpleNamespace(vllm_enable_sleep_mode=False)
+    trainer.state = SimpleNamespace(global_step=0)
+    trainer._last_loaded_step = 0
+    trainer.processing_class = object()
+    trainer.llm = DummyLLM()
+    trainer.generation_config = object()
+    image = object()
+    prompts = [
+        [{"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": "Image"}]}],
+        [{"role": "user", "content": [{"type": "text", "text": "Text"}]}],
+    ]
+
+    trainer._generate_vllm_colocate(prompts, images=[image, None])
+
+    assert trainer.llm.inputs[0] == {"prompt": "image text", "multi_modal_data": {"image": image}}
+    assert trainer.llm.inputs[1] == "text"
+
+
 class TestOnlineDPOTrainer(TrlTestCase):
     def setup_method(self):
         self.model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"

trl/experimental/online_dpo/online_dpo_trainer.py

@@ -653,13 +653,8 @@ def _generate_vllm_server(self, prompts, images=None):
         # Gather all prompts to main process
         all_prompts = gather_object(prompts_text)
         if has_images:
-            # The server can't take images alongside text prompts, so multimodal prompts are sent as messages, with
-            # the images inlined in place of their placeholders.
-            messages = [
-                prepare_multimodal_messages(prompt, images=[image] if image is not None else None)
-                for prompt, image in zip(prompts, images, strict=True)
-            ]
-            all_messages = gather_object(messages)
+            # The server can't take images alongside text prompts, so multimodal prompts are sent as messages.
+            all_messages = gather_object(prompts)
 
         if self.accelerator.is_main_process:
             sampling_kwargs = {
@@ -963,7 +958,7 @@ def _generate(self, model, prompts, images=None):
         # Prepare kwargs for processing class
         kwargs = {}
         if images is not None:
-            kwargs = {"images": [[img] for img in images]}
+            kwargs = {"images": images}
 
         # Process inputs using the processing class (handles both VLM and LLM)
         prompt_inputs = self.processing_class(
@@ -1122,40 +1117,44 @@ def training_step(
         batch_size = len(prompts)
 
         # Handle images for VLM support
-        has_images = "image" in inputs
         images = None
-        if has_images:
+        image_lists = None
+        if "image" in inputs:
             images = inputs["image"]
-            # Convert conversational prompts to include image tokens
-            for prompt in prompts:
-                if isinstance(prompt, list):
-                    for message in prompt:
-                        if not isinstance(message, dict):
-                            continue
-                        content = message.get("content")
-                        role = message.get("role")
-                        if isinstance(content, str):
-                            if role == "user":
-                                message["content"] = [{"type": "image"}, {"type": "text", "text": content}]
-                            elif role == "system":
-                                message["content"] = [{"type": "text", "text": content}]
+            image_lists = [[image] if image is not None else [] for image in images]
+            if all(image_list == [] for image_list in image_lists):
+                images = None
+                image_lists = None
+
+        if images is not None:
+            if not is_conversational({"prompt": prompts[0]}):
+                raise ValueError(
+                    "Multimodal training requires conversational prompts. It looks like the dataset contains "
+                    "non-conversational inputs, likely because a chat template was applied before passing the dataset "
+                    "to the trainer. Please provide the raw conversational prompts and let the trainer apply the chat "
+                    "template internally."
+                )
+            prompts = [
+                prepare_multimodal_messages(prompt, images=image_list)
+                for prompt, image_list in zip(prompts, image_lists, strict=True)
+            ]
 
         if self.args.use_vllm:
             prompt_ids, prompt_mask, completion_ids, completion_mask = self._generate_vllm(prompts, images)
         else:
-            prompt_ids, prompt_mask, completion_ids, completion_mask = self._generate(model, prompts, images)
+            prompt_ids, prompt_mask, completion_ids, completion_mask = self._generate(model, prompts, image_lists)
 
         contain_eos_token = torch.any(completion_ids == self._tokenizer.eos_token_id, dim=-1)
 
         # Extract vision inputs if available for VLM support
         vision_inputs = None
-        if has_images and self.is_vision_model and not self.args.use_vllm:
+        if images is not None and self.is_vision_model and not self.args.use_vllm:
             # For vision models with transformers generation, we need to prepare vision inputs
             # Process the images to get vision inputs that can be passed through the forward pass
             vision_inputs = {}
-            kwargs = {"images": [[img] for img in images]}
+            kwargs = {"images": image_lists}
             processed = self.processing_class(
-                text=[""] * len(images),  # Dummy text for vision processing
+                text=[""] * len(image_lists),  # Dummy text for vision processing
                 return_tensors="pt",
                 **kwargs,
             )