Support max_length in DPO VLM training

tests/test_dpo_trainer.py

@@ -1212,6 +1212,26 @@ def test_train_vlm_text_only_data(self, model_id, dataset_config):
             else:
                 assert not torch.allclose(param, new_param, rtol=1e-12, atol=1e-12), f"Param {n} is not updated"
 
+    @require_vision
+    def test_train_vlm_with_max_length(self):
+        # Regression test for #5283: mm_token_type_ids must be truncated alongside input_ids when max_length is set,
+        # otherwise a shape mismatch crashes the model forward pass.
+        # max_length=37 truncates 1 completion token (total_len=38) while keeping all image tokens (prompt_len=34) safe.
+        dataset = load_dataset("trl-internal-testing/zen-image", "conversational_preference", split="train")
+        training_args = DPOConfig(
+            output_dir=self.tmp_dir,
+            max_length=37,  # total_len=38, prompt_len=34 — truncates completion, not image tokens
+            per_device_train_batch_size=2,
+            report_to="none",
+        )
+        trainer = DPOTrainer(
+            model="trl-internal-testing/tiny-Qwen2_5_VLForConditionalGeneration",
+            args=training_args,
+            train_dataset=dataset,
+        )
+        trainer.train()
+        assert trainer.state.log_history[-1]["train_loss"] is not None
+
     @require_peft
     @require_bitsandbytes
     def test_peft_with_quantization(self):

trl/trainer/dpo_trainer.py

@@ -1015,6 +1015,10 @@ def compute_ref_log_probs(self, inputs):
         ):
             if key in inputs:
                 model_kwargs[key] = inputs[key]
+        # token_type_ids and mm_token_type_ids are sequence-length-aligned; truncate to match input_ids
+        for key in ("token_type_ids", "mm_token_type_ids"):
+            if key in model_kwargs:
+                model_kwargs[key] = model_kwargs[key][:, : input_ids.shape[1]]
 
         with torch.no_grad(), disable_gradient_checkpointing(self.model, self.args.gradient_checkpointing_kwargs):
             if is_peft_model(self.model) and self.ref_model is None:
@@ -1143,6 +1147,10 @@ def _compute_loss(self, model, inputs, return_outputs):
         ):
             if key in inputs:
                 model_kwargs[key] = inputs[key]
+        # token_type_ids and mm_token_type_ids are sequence-length-aligned; truncate to match input_ids
+        for key in ("token_type_ids", "mm_token_type_ids"):
+            if key in model_kwargs:
+                model_kwargs[key] = model_kwargs[key][:, : input_ids.shape[1]]
 
         outputs = model(**model_kwargs)
         shift_logits = outputs.logits[..., :-1, :].contiguous()