Skip to content

Fix bug with VLM processors in prompt-completion completion text-only training #4553

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
qgallouedec merged 3 commits into from
Nov 20, 2025
Merged

Fix bug with VLM processors in prompt-completion completion text-only training #4553

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
0265d38
fix(SFTTrainer): Support VLM processors in completion-only text training
kschwethelm Nov 20, 2025
3cd4938
Added pytests for vlm text only data prompt completion
kschwethelm Nov 20, 2025
93db234
parametrize
qgallouedec Nov 20, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions tests/test_sft_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1532,10 +1532,14 @@ def test_train_vlm_gemma_3n(self):
"trl-internal-testing/tiny-Qwen2_5_VLForConditionalGeneration",
],
)
@pytest.mark.parametrize(
"dataset_config",
["conversational_language_modeling", "conversational_prompt_completion", "standard_prompt_completion"],
)
@require_vision
def test_train_vlm_text_only_data(self, model_id):
def test_train_vlm_text_only_data(self, model_id, dataset_config):
# Get the dataset
dataset = load_dataset("trl-internal-testing/zen", "conversational_language_modeling", split="train")
dataset = load_dataset("trl-internal-testing/zen", dataset_config, split="train")

# Initialize the trainer
training_args = SFTConfig(output_dir=self.tmp_dir, report_to="none")
Expand Down
11 changes: 10 additions & 1 deletion trl/trainer/sft_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -935,9 +935,10 @@ def add_eos(example, eos_token):
example["completion"] = example["completion"] + eos_token
return example

eos_token = processing_class.tokenizer.eos_token if self._is_vlm else processing_class.eos_token
dataset = dataset.map(
add_eos,
fn_kwargs={"eos_token": processing_class.eos_token},
fn_kwargs={"eos_token": eos_token},
remove_columns="messages" if "messages" in column_names else None, # renamed to "text"
**map_kwargs,
)
Expand Down Expand Up @@ -988,6 +989,14 @@ def tokenize_fn(example, processing_class, dataset_text_field, assistant_only_lo
prompt_completion_ids = processing_class(text=example["prompt"] + example["completion"])[
"input_ids"
]
# Fix transformers inconsistency: for VLMs, processing_class returns lists of lists
# even for single examples, while for LLMs it returns lists of ints.
prompt_ids = prompt_ids[0] if isinstance(prompt_ids[0], list) else prompt_ids
prompt_completion_ids = (
prompt_completion_ids[0]
if isinstance(prompt_completion_ids[0], list)
else prompt_completion_ids
)

# Check if the tokenized prompt starts with the tokenized prompt+completion
if not prompt_completion_ids[: len(prompt_ids)] == prompt_ids:
Expand Down
Loading