From 4066f7fa374d6c51a420b630f207110962be0bd4 Mon Sep 17 00:00:00 2001 From: Daoyuan Li <94409450+DaoyuanLi2816@users.noreply.github.com> Date: Fri, 28 Aug 2026 22:32:09 -0700 Subject: [PATCH 1/3] Fix mixed-image Online DPO server batches --- tests/experimental/test_online_dpo_trainer.py | 46 ++++++++++++++++++- .../online_dpo/online_dpo_trainer.py | 19 ++++++-- 2 files changed, 60 insertions(+), 5 deletions(-) diff --git a/tests/experimental/test_online_dpo_trainer.py b/tests/experimental/test_online_dpo_trainer.py index 7e8a18d4d11..65bb6429c02 100644 --- a/tests/experimental/test_online_dpo_trainer.py +++ b/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,47 @@ from transformers import AutoModelForImageTextToText, AutoProcessor +def test_generate_vllm_server_with_mixed_image_batch(monkeypatch): + class DummyProcessor: + def __call__(self, **kwargs): + return {"input_ids": torch.tensor([[1], [2]])} + + class DummyClient: + def chat(self, messages, **kwargs): + self.messages = messages + return {"completion_ids": [[3], [4], [5], [6]]} + + monkeypatch.setattr(online_dpo_trainer, "apply_chat_template", lambda example, processor: {"prompt": "prompt"}) + 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.processing_class = DummyProcessor() + trainer.accelerator = SimpleNamespace(is_main_process=True, process_index=0) + 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) + trainer.vllm_client = DummyClient() + + image = object() + prompts = [ + [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "Describe it"}]}], + [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "Say hello"}]}], + ] + + trainer._generate_vllm_server(prompts, images=[image, None]) + + 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"}] + + class TestOnlineDPOTrainer(TrlTestCase): def setup_method(self): self.model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5" diff --git a/trl/experimental/online_dpo/online_dpo_trainer.py b/trl/experimental/online_dpo/online_dpo_trainer.py index 134670a49e8..3c0ea47808c 100644 --- a/trl/experimental/online_dpo/online_dpo_trainer.py +++ b/trl/experimental/online_dpo/online_dpo_trainer.py @@ -655,10 +655,21 @@ def _generate_vllm_server(self, prompts, images=None): 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) - ] + messages = [] + for prompt, image in zip(prompts, images, strict=True): + if image is None: + prompt = [ + { + **message, + "content": [ + part for part in message["content"] if part.get("type") != "image" or "image" in part + ], + } + if isinstance(message.get("content"), list) + else message + for message in prompt + ] + messages.append(prepare_multimodal_messages(prompt, images=[image] if image is not None else None)) all_messages = gather_object(messages) if self.accelerator.is_main_process: From 70f0231d229664abf6ccab8c9c72e0b0bf9b05ac Mon Sep 17 00:00:00 2001 From: Daoyuan Li <94409450+DaoyuanLi2816@users.noreply.github.com> Date: Sat, 29 Aug 2026 10:26:31 -0700 Subject: [PATCH 2/3] Align mixed-batch prompt tokenization --- tests/experimental/test_online_dpo_trainer.py | 10 +++++- .../online_dpo/online_dpo_trainer.py | 35 +++++++++++-------- 2 files changed, 29 insertions(+), 16 deletions(-) diff --git a/tests/experimental/test_online_dpo_trainer.py b/tests/experimental/test_online_dpo_trainer.py index 65bb6429c02..0a8bbd328cc 100644 --- a/tests/experimental/test_online_dpo_trainer.py +++ b/tests/experimental/test_online_dpo_trainer.py @@ -38,6 +38,7 @@ def test_generate_vllm_server_with_mixed_image_batch(monkeypatch): class DummyProcessor: def __call__(self, **kwargs): + self.text = kwargs["text"] return {"input_ids": torch.tensor([[1], [2]])} class DummyClient: @@ -45,7 +46,13 @@ def chat(self, messages, **kwargs): self.messages = messages return {"completion_ids": [[3], [4], [5], [6]]} - monkeypatch.setattr(online_dpo_trainer, "apply_chat_template", lambda example, processor: {"prompt": "prompt"}) + 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) @@ -74,6 +81,7 @@ def chat(self, messages, **kwargs): 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.processing_class.text == ["image text", "text"] class TestOnlineDPOTrainer(TrlTestCase): diff --git a/trl/experimental/online_dpo/online_dpo_trainer.py b/trl/experimental/online_dpo/online_dpo_trainer.py index 3c0ea47808c..352b2149675 100644 --- a/trl/experimental/online_dpo/online_dpo_trainer.py +++ b/trl/experimental/online_dpo/online_dpo_trainer.py @@ -645,6 +645,22 @@ def _generate_vllm_server(self, prompts, images=None): self._move_model_to_vllm() self._last_loaded_step = self.state.global_step + if has_images: + prompts = [ + [ + { + **message, + "content": [ + part for part in message["content"] if part.get("type") != "image" or "image" in part + ], + } + if image is None and isinstance(message.get("content"), list) + else message + for message in prompt + ] + for prompt, image in zip(prompts, images, strict=True) + ] + # Apply chat template if conversational if is_conversational({"prompt": prompts[0]}): prompts_text = [apply_chat_template({"prompt": p}, self.processing_class)["prompt"] for p in prompts] @@ -655,21 +671,10 @@ def _generate_vllm_server(self, prompts, images=None): 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 = [] - for prompt, image in zip(prompts, images, strict=True): - if image is None: - prompt = [ - { - **message, - "content": [ - part for part in message["content"] if part.get("type") != "image" or "image" in part - ], - } - if isinstance(message.get("content"), list) - else message - for message in prompt - ] - messages.append(prepare_multimodal_messages(prompt, images=[image] if image is not None else None)) + 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) if self.accelerator.is_main_process: From a9d918c601333547beb4d24ab69d485de8697999 Mon Sep 17 00:00:00 2001 From: Daoyuan Li <94409450+DaoyuanLi2816@users.noreply.github.com> Date: Thu, 3 Sep 2026 22:02:20 -0700 Subject: [PATCH 3/3] Fix mixed-image generation across Online DPO modes --- tests/experimental/test_online_dpo_trainer.py | 113 ++++++++++++++++-- .../online_dpo/online_dpo_trainer.py | 69 ++++------- 2 files changed, 130 insertions(+), 52 deletions(-) diff --git a/tests/experimental/test_online_dpo_trainer.py b/tests/experimental/test_online_dpo_trainer.py index 0a8bbd328cc..20cc4e525d7 100644 --- a/tests/experimental/test_online_dpo_trainer.py +++ b/tests/experimental/test_online_dpo_trainer.py @@ -35,16 +35,19 @@ from transformers import AutoModelForImageTextToText, AutoProcessor -def test_generate_vllm_server_with_mixed_image_batch(monkeypatch): +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]])} + return {"input_ids": torch.tensor([[1], [2], [3]])} class DummyClient: def chat(self, messages, **kwargs): self.messages = messages - return {"completion_ids": [[3], [4], [5], [6]]} + return {"completion_ids": [[4], [5], [6], [7], [8], [9]]} monkeypatch.setattr( online_dpo_trainer, @@ -59,8 +62,9 @@ def chat(self, messages, **kwargs): 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) + 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 @@ -68,20 +72,111 @@ def chat(self, messages, **kwargs): trainer.top_k = None trainer.min_p = None trainer.generation_config = SimpleNamespace(max_tokens=16) - trainer.args = SimpleNamespace(generation_kwargs=None) + 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": [{"type": "image"}, {"type": "text", "text": "Describe it"}]}], - [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "Say hello"}]}], + [{"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 - trainer._generate_vllm_server(prompts, images=[image, None]) + 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.processing_class.text == ["image text", "text"] + 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): diff --git a/trl/experimental/online_dpo/online_dpo_trainer.py b/trl/experimental/online_dpo/online_dpo_trainer.py index 352b2149675..8b706d925b4 100644 --- a/trl/experimental/online_dpo/online_dpo_trainer.py +++ b/trl/experimental/online_dpo/online_dpo_trainer.py @@ -645,22 +645,6 @@ def _generate_vllm_server(self, prompts, images=None): self._move_model_to_vllm() self._last_loaded_step = self.state.global_step - if has_images: - prompts = [ - [ - { - **message, - "content": [ - part for part in message["content"] if part.get("type") != "image" or "image" in part - ], - } - if image is None and isinstance(message.get("content"), list) - else message - for message in prompt - ] - for prompt, image in zip(prompts, images, strict=True) - ] - # Apply chat template if conversational if is_conversational({"prompt": prompts[0]}): prompts_text = [apply_chat_template({"prompt": p}, self.processing_class)["prompt"] for p in prompts] @@ -669,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 = { @@ -979,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( @@ -1138,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, )