Support max_length in DPO VLM training

[albertvillanova]

Support max_length in DPO VLM training.

• Truncate sequence-aligned side-inputs (token_type_ids, mm_token_type_ids) with input_ids in DPO VLM training

Fix #5283.

This PR addresses a regression affecting vision-language model (VLM) training when using sequence truncation. The main fix ensures that auxiliary token fields (mm_token_type_ids and token_type_ids) are truncated in sync with input_ids, preventing shape mismatches and crashes during the model's forward pass. Additionally, a regression test is added to verify this behavior.

Changes

Bug fix for sequence truncation in VLMs:

• Ensured that token_type_ids and mm_token_type_ids are truncated to match the length of input_ids in both compute_ref_log_probs and _compute_loss methods of DPOTrainer, preventing shape mismatch errors during training.
  • Note that pixel_values, image_grid_thw, image_sizes, and pixel_attention_mask are patch-level or image-level tensors and should not be truncated.

Testing improvements:

• Added a regression test test_train_vlm_with_max_length in tests/test_dpo_trainer.py to verify that truncation with max_length does not crash the model and that image tokens are handled correctly.

Follow-up

If this approach is approved, I will implement a similar fix for other trainers.

Related

See related discussion in: #5279 (comment)

---

Note

Medium Risk
Changes DPOTrainer truncation behavior for sequence-aligned auxiliary tensors in VLM training; incorrect handling could affect training correctness or cause new shape/attention issues when max_length/truncation_mode are used.

Overview
Fixes a VLM regression when training DPOTrainer with max_length by truncating sequence-aligned side inputs (currently token_type_ids and mm_token_type_ids) in lockstep with input_ids/masks.

Extends _truncate_inputs to accept and return additional tensors, and wires this through both compute_ref_log_probs and _compute_loss while leaving image/patch-level inputs (e.g. pixel_values, image_grid_thw) untouched.

Adds a require_vision regression test test_train_vlm_with_max_length to ensure truncated VLM batches no longer crash due to shape mismatches.

^{Written by Cursor Bugbot for commit 38b4bc6. This will update automatically on new commits. Configure here.}

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

[qgallouedec]

why not having mm_token_type_ids in extra?

[albertvillanova]

The tokens in extra can be truncated both with "keep_start" and "keep_end", and I think it is semantically wrong to use "keep_end" in VLM mm_token_type_ids, but I'm addressing that in a following PR:

• Prevent corruption of DPO VLM training if "keep_end" truncation_mode #5286

So, let's treat mm_token_type_ids and token_type_ids symmetrically to be internally consistent, and leave the semantical correction to the other PR.

[qgallouedec]

ok, I think you should align compute_ref_log_probs with compute_loss, ie having mm_token_type_ids in extra in both cases

[albertvillanova]

Sure, thanks! 😅

[albertvillanova]

Done!