Fix VLM model + text-only dataset ValueError in TRL 0.22.x (#4004)

TRL 0.22.x checks _is_vlm (model type) instead of _is_vision_dataset
(dataset content, added in 0.25.1+) in _set_signature_columns_if_needed.
When _is_vlm=True (e.g. Gemma3), signature columns are set to vision-only
["messages","prompt","completion","images"], which has zero overlap with
tokenized text columns [input_ids, labels, attention_mask, ...], causing
a ValueError.

Fix: expand the VLM branch signature columns to include both vision and
text column names. Extra columns not present in the dataset are harmlessly
ignored by _remove_unused_columns (it only raises when zero columns match).

Co-authored-by: Daniel Hanchen <danielhanchen@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-02-09 06:24:58 -08:00 committed by GitHub
commit 52e35bbfd7

View file

@ -1276,6 +1276,22 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
_vlm_check_original, _vlm_check_patched
)
# Fix TRL 0.22.x: VLM models with text-only datasets.
# TRL 0.22.x checks _is_vlm (model type) not _is_vision_dataset (dataset
# content, added in 0.25.1+). When _is_vlm=True, signature columns are
# vision-only ["messages","prompt","completion","images"], which have zero
# overlap with tokenized text columns. Fix: merge both column sets into the
# VLM branch. Extra columns not in the dataset are harmlessly ignored by
# _remove_unused_columns (it only raises when zero columns match).
_sig_vlm_old = (
'self._signature_columns = ["messages", "prompt", "completion", "images"]'
)
_sig_vlm_new = (
'self._signature_columns = ["messages", "prompt", "completion", "images",'
' "input_ids", "labels", "attention_mask", "seq_lengths", "completion_mask", "assistant_masks"]'
)
RLTrainer_source = RLTrainer_source.replace(_sig_vlm_old, _sig_vlm_new)
# Remove multiple doc strings
if __RLConfig_doc__ != "" and RLTrainer_source.count(__RLTrainer_doc__) == 2:
RLTrainer_source = RLTrainer_source.replace(__RLTrainer_doc__, "", 1)