diff --git a/tests/experimental/test_online_dpo_trainer.py b/tests/experimental/test_online_dpo_trainer.py index 7e8a18d4d11..20cc4e525d7 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,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" diff --git a/trl/experimental/online_dpo/online_dpo_trainer.py b/trl/experimental/online_dpo/online_dpo_trainer.py index 134670a49e8..8b706d925b4 100644 --- a/trl/experimental/online_dpo/online_dpo_trainer.py +++ b/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, )