huggingface / huggingface/trl

GRPO + vLLM corrupts SmolVLM multimodal prompts from pre-expanded image tokens

Open
#6,294 4 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
19.3k
Forks
3k
Avg merge
1d 20h
Merged PRs (30d)
194

Description

### Reproduction

When using `GRPOTrainer` with `use_vllm=True, vllm_mode="colocate"`, and `HuggingFaceTB/SmolVLM-Instruct`, the current TRL path sends processor-expanded multimodal `prompt_token_ids` to vLLM, which leads to the broken completions.

The minimal diagnostic showing the prompt mismatch:
```python
from datasets import load_dataset
from transformers import AutoProcessor

MODEL_ID = "HuggingFaceTB/SmolVLM-Instruct"
DATASET_ID = "leonardPKU/GEOQA_R1V_Train_8K"

def main():
processor = AutoProcessor.from_pretrained(MODEL_ID)
example = load_dataset(DATASET_ID, split="train[:1]")[0]
image = example["image"].convert("RGB")

prompt = [
{
"role": "user",
"content": [
{"type": "image", "image": image},
{"type": "text", "text": example["problem"]},
],
}
]

expanded = processor.apply_chat_template(
prompt,
add_generation_prompt=True,
tokenize=True,
return_dict=True,
)["input_ids"][0]

prompt_text = processor.apply_chat_template(
prompt,
add_generation_prompt=True,
tokenize=False,
)
unexpanded = processor.tokenizer(prompt_text, add_special_tokens=False)["input_ids"]

image_token = getattr(processor.tokenizer, "image_token", "")
image_token_id = processor.tokenizer.convert_tokens_to_ids(image_token)

print(f"model: {MODEL_ID}")
print(f"dataset: {DATASET_ID}")
print(f"problem: {example['problem']}")
print()
print(f"image token: {image_token!r}")
print(f"image token id: {image_token_id}")
print(f"processor-expanded prompt length: {len(expanded)}")
print(f"tokenizer-only prompt length: {len(unexpanded)}")
print(f"processor-expanded image-token count: {expanded.count(image_token_id)}")
print(f"tokenizer-only image-token count: {unexpanded.count(image_token_id)}")
print()
print("First 40 processor-expanded IDs:")
print(expanded[:40])
print()
print("First 40 tokenizer-only IDs:")
print(unexpanded[:40])
print()
print("Prompt text prefix:")
print(prompt_text[:500])

if __name__ == "__main__":
main()
```
Output:
```text
>image token: ''
>image token id: 49153
>processor-expanded prompt length: 855
>tokenizer-only prompt length: 38
>processor-expanded image-token count: 729
>tokenizer-only image-token count: 1
```

The GRPO example I run as A/B test with my monkey-patch fix (**2 .csv artifacts from this run to confirm broken completions you can see in attached files**):
[patched_rollouts.csv](https://github.com/user-attachments/files/29675599/patched_rollouts.csv)
[current_method_rollouts.csv](https://github.com/user-attachments/files/29675598/current_method_rollouts.csv)
```python
"""
Debug GRPO run: SmolVLM-Instruct + LoRA (r=4, alpha=8, q/k projections) on GEOQA_R1V_Train_8K.

Runs two vLLM (colocate mode, single GPU) GRPO jobs and dumps every rollout to
separate CSV files for inspection:
- grpo_smolvlm_debug/current_method_rollouts.csv: stock TRL prompt-token path
- grpo_smolvlm_debug/patched_rollouts.csv: unexpanded vLLM prompt-token path

Run:
.venv/bin/python grpo_smolvlm_geoqa_debug.py
"""

import argparse
import csv
import gc
import os
import re
import subprocess
import sys

import torch
from datasets import load_dataset
from peft import LoraConfig

from trl import GRPOConfig, GRPOTrainer
from trl.extras.profiling import profiling_context

SYSTEM_PROMPT = """
Answer the question by briefly explaining the reasoning behind your answer.
Return the final answer as a single number followed immediately by the ° symbol.
"""

OUTPUT_DIR = "grpo_smolvlm_debug"
CURRENT_METHOD_CSV_PATH = os.path.join(OUTPUT_DIR, "current_method_rollouts.csv")
PATCHED_CSV_PATH = os.path.join(OUTPUT_DIR, "patched_rollouts.csv")

# Last number immediately followed by the ° symbol
NUMBER_DEG_RE = re.compile(r"(\d+(?:\.\d+)?)\s*°")

def normalize_solution(solution):
solution = str(solution).replace("", "").replace("", "").strip()
if solution and not solution.endswith("°"):
solution = f"{solution}°"
return solution

def parse_answer(text):
matches = NUMBER_DEG_RE.findall(text)
return f"{matches[-1]}°" if matches else ""

def shorten_cell(value, width):
value = " ".join(str(value).split())
return value if len(value) <= width else value[: width - 1] + "…"

def print_rollouts_table(run_name, step, rows):
headers = ["idx", "gold", "parsed", "reward", "problem", "completion"]
widths = [4, 10, 10, 8, 44, 80]
print(f"\n=== {run_name} step {step} rollouts ===", flush=True)
print(" | ".join(header.ljust(width) for header, width in zip(headers, widths)), flush=True)
print("-+-".join("-" * width for width in widths), flush=True)
for idx, row in enumerate(rows):
_, _, question, gold, text, answered, reward = row
values = [idx, gold, answered, f"{reward:.1f}", question, text]
print(" | ".join(shorten_cell(value, width).ljust(width) for value, width in zip(values, widths)), flush=True)

def make_exact_answer_reward(csv_path, run_name):
def exact_answer_reward(prompts, completions, problem, solution, trainer_state, **kwargs):
rewards = []
rows = []
for question, completion, sol in zip(problem, completions, solution):
text = completion[0]["content"] if isinstance(completion, list) else completion
if isinstance(text, list): # content blocks -> plain text
text = "".join(part.get("text", "") for part in text)
gold = normalize_solution(sol)
answered = parse_answer(text)
gold_match = NUMBER_DEG_RE.findall(gold)
if gold_match and answered:
# numeric compare so "145°" == "145.0°"
reward = 1.0 if float(answered.rstrip("°")) == float(gold_match[-1]) else 0.0
else:
reward = 1.0 if answered == gold else 0.0
rewards.append(reward)
rows.append([run_name, trainer_state.global_step, question, gold, text, answered, reward])

print_rollouts_table(run_name, trainer_state.global_step, rows)

write_header = not os.path.exists(csv_path)
with open(csv_path, "a", newline="") as f:
writer = csv.writer(f)
if write_header:
writer.writerow(["run", "step", "problem", "gold_solution", "completion", "parsed_answer", "reward"])
writer.writerows(rows)
return rewards

return exact_answer_reward

class PatchedSmolVLMGRPOTrainer(GRPOTrainer):
"""
Integrated workaround for vLLM (0.13-0.23) x Idefics3/SmolVLM prompt corruption.

TRL's colocate path sends HF-pre-expanded `prompt_token_ids` to vLLM, but for token-id
prompts vLLM always re-applies its prompt updates (`is_update_applied=False` in
`_apply_hf_processor_main`). Idefics3's replacement target is the bare `` token, so
vLLM replaces the first of the 729 already-expanded `` tokens with a fresh full
expansion (855 -> 1672 tokens) and the image features end up misaligned -> garbage rollouts.

Fix: keep the expanded ids for the training forward pass, but feed vLLM unexpanded prompt ids
produced from chat-template text tokenized without the image processor.
"""

def _tokenize_prompts(self, prompts: list):
prompt_ids, images, multimodal_fields = super()._tokenize_prompts(prompts)
self._vllm_prompt_ids = prompt_ids
if self.use_vllm and images is not None:
texts = self.processing_class.apply_chat_template(
conversation=prompts,
chat_template=self.chat_template,
add_generation_prompt=True,
tokenize=False,
**self.chat_template_kwargs,
)
self._vllm_prompt_ids = self.processing_class.tokenizer(texts, add_special_tokens=False)["input_ids"]
return prompt_ids, images, multimodal_fields

def _generate_single_turn(self, prompt_ids, images, multimodal_fields):
if not self.use_vllm:
return super()._generate_single_turn(prompt_ids, images, multimodal_fields)

mode = "train" if self.model.training else "eval"

# Sync weights if training step changed
if self.state.global_step != self._last_loaded_step:
with profiling_context(self, "sync_weights"):
self.vllm_generation.sync_weights()
self._last_loaded_step = self.state.global_step

# Generate using vLLM with unexpanded prompt IDs. The expanded IDs remain the source of truth
# for the training forward pass in the parent GRPO pipeline.
num_generations = self.num_generations if mode == "train" else self.num_generations_eval
_, completion_ids, logprobs, _ = self.vllm_generation.generate(
prompts=self._vllm_prompt_ids,
images=images,
num_generations=num_generations,
profiler=profiling_context(self, "vLLM.generate"),
)
# vLLM returns per-token top-k logprobs; keep only the top-1 (sampled token) logprob
logprobs = [[lp[0] for lp in seq] for seq in logprobs]
return completion_ids, logprobs

def build_dataset():
dataset = load_dataset("leonardPKU/GEOQA_R1V_Train_8K", split="train")
dataset = dataset.select(range(32)) # 5 steps x 4 rollouts -> 20 samples needed

def make_conversation(example):
return {
"prompt": [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": example["problem"]},
]
}

dataset = dataset.map(make_conversation)

def convert_to_rgb(example):
image = example["image"]
if image.mode != "RGB":
example["image"] = image.convert("RGB")
return example

return dataset.map(convert_to_rgb)

def build_training_args(output_dir):
return GRPOConfig(
output_dir=output_dir,
learning_rate=1e-5,
per_device_train_batch_size=2,
gradient_accumulation_steps=2,
num_generations=4,
max_completion_length=256,
max_steps=5,
bf16=True,
gradient_checkpointing=True,
logging_steps=1,
save_strategy="no",
report_to="none",
use_vllm=True,
vllm_mode="colocate",
vllm_gpu_memory_utilization=0.4,
vllm_max_model_length=4096,
model_init_kwargs={"dtype": torch.bfloat16},
seed=42,
)

def run_experiment(run_name, trainer_cls, csv_path, dataset):
print(f"Running {run_name}; writing rollouts to {csv_path}")
if os.path.exists(csv_path):
os.remove(csv_path)

peft_config = LoraConfig(
r=4,
lora_alpha=8,
target_modules=["q_proj", "k_proj"],
task_type="CAUSAL_LM",
)

trainer = trainer_cls(
model="HuggingFaceTB/SmolVLM-Instruct",
args=build_training_args(os.path.join(OUTPUT_DIR, run_name)),
reward_funcs=make_exact_answer_reward(csv_path, run_name),
train_dataset=dataset,
peft_config=peft_config,
)
trainer.train()
trainer.accelerator.free_memory()
del trainer
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()

def run_requested_experiment(run_name):
os.makedirs(OUTPUT_DIR, exist_ok=True)
dataset = build_dataset()
if run_name == "current_method":
run_experiment(run_name, GRPOTrainer, CURRENT_METHOD_CSV_PATH, dataset)
elif run_name == "patched":
run_experiment(run_name, PatchedSmolVLMGRPOTrainer, PATCHED_CSV_PATH, dataset)
else:
raise ValueError(f"Unknown run: {run_name}")

def main():
parser = argparse.ArgumentParser()
parser.add_argument("--run", choices=["current_method", "patched"])
args = parser.parse_args()

if args.run is not None:
run_requested_experiment(args.run)
return

os.makedirs(OUTPUT_DIR, exist_ok=True)
for run_name in ["current_method", "patched"]:
subprocess.run([sys.executable, __file__, "--run", run_name], check=True)

print(f"Wrote {CURRENT_METHOD_CSV_PATH}")
print(f"Wrote {PATCHED_CSV_PATH}")

if __name__ == "__main__":
main()
```
A 5-step A/B test summarization:
- stock `GRPOTrainer`: 20 rollouts, 10 non-empty completions; the non-empty ones were mostly nonsensical; 1 parsed answer
- patched subclass: 20 rollouts, 20 non-empty completions; outputs were visibly more coherent; 4 parsed answers

### Expected behavior

For multimodal vLLM generation, TRL should pass unexpanded tokenizer-only prompt IDs to vLLM, together with multi_modal_data, and let vLLM apply its own image prompt expansion. The processor-expanded IDs should still be used for the training forward pass/logprob computation.

### Suggested fix
```python
# In GRPOTrainer._tokenize_prompts(...)
prompt_ids, images, multimodal_fields = ...
vllm_prompt_ids = prompt_ids

if self.use_vllm and images is not None:
texts = self.processing_class.apply_chat_template(
conversation=prompts,
tools=self.tools or None,
chat_template=self.chat_template,
add_generation_prompt=True,
tokenize=False,
**self.chat_template_kwargs,
)
vllm_prompt_ids = self.processing_class.tokenizer(texts, add_special_tokens=False)["input_ids"]

return prompt_ids, images, multimodal_fields, vllm_prompt_ids
```

Then at the vLLM boundary:
```python
# In GRPOTrainer._generate(...)
prompt_ids, images, multimodal_fields, vllm_prompt_ids = self._tokenize_prompts(prompts)
completion_ids, logprobs = self._generate_single_turn(
prompt_ids,
images,
multimodal_fields,
vllm_prompt_ids,
)

# In GRPOTrainer._generate_single_turn(...)
_, completion_ids, logprobs, _ = self.vllm_generation.generate(
prompts=vllm_prompt_ids if self.use_vllm else prompt_ids,
images=images,
num_generations=num_generations,
profiler=profiling_context(self, "vLLM.generate"),
)
```

I’m happy to open a PR if this approach looks reasonable. Since this logic is duplicated, the same change likely needs to be mirrored in all trainers with on-policy generation as well, not only `GRPOTrainer`.

### Update — also affects vLLM server mode

This also affects `vllm_mode="server"`, where it surfaces as a hard failure rather than corrupted rollouts. Noting it here since it's the same root cause.

Running the reproducer above against a live vLLM server on current `main` (`6630e17a`) with `HuggingFaceTB/SmolVLM-Instruct`:

```
ValueError: Found 9 runs of image tokens in the prompt but 1 images were processed.
The prompt must contain one run of image tokens per image.
```

Server mode no longer sends images to vLLM. It processes them separately via `VLLMClient.image_features` and then locates them with `VLLMGeneration._place_features`, which scans the prompt for runs of `image_token_id` and requires exactly one run per image. Idefics3 expansion interleaves `` runs with structural tokens (``, ``), so the 729 image tokens form **9 runs of 81** — never a single run — and the guard raises at step 0.

The server's own `image_features` reports the image as one contiguous span (`{'offset': 3, 'length': 818}`), which is what an un-expanded prompt yields and an expanded one never does. Feeding un-expanded IDs, the same run completes: 8 rollouts, 0 empty, all coherent, one exact answer match.

So the un-expanded prompt IDs are needed in both vLLM modes, for different reasons — colocate re-expands the placeholders, server can't locate them. #6718 covers both.

Verified end to end on SmolVLM only. Qwen-style VLMs expand into a single contiguous `<|image_pad|>` run, so they already satisfy `_place_features` today.

### System Info

```text
- Platform: Linux-6.17.0-35-generic-x86_64-with-glibc2.39
- Python version: 3.13.5
- TRL version: 1.5.0.dev0+b73bc7b
- PyTorch version: 2.11.0
- accelerator(s): NVIDIA GeForce RTX 5070 Ti
- Transformers version: 5.13.0
- Accelerate version: 1.13.0
- Accelerate config:
- compute_environment: LOCAL_MACHINE
- distributed_type: NO
- mixed_precision: bf16
- use_cpu: False
- debug: False
- num_processes: 1
- machine_rank: 0
- num_machines: 1
- gpu_ids: 1
- rdzv_backend: static
- same_network: True
- main_training_function: main
- enable_cpu_affinity: True
- downcast_bf16: no
- tpu_use_cluster: False
- tpu_use_sudo: False
- tpu_env: []
- Datasets version: 4.8.5
- HF Hub version: 1.22.0
- bitsandbytes version: 0.49.2
- DeepSpeed version: 0.18.9
- Liger-Kernel version: 0.8.0
- PEFT version: 0.18.1
- vLLM version: 0.23.0
```

### Checklist

- [x] I have checked that my issue isn't already filed (see [open issues](https://github.com/huggingface/trl/issues?q=is%3Aissue))
- [x] I have included my system information
- [x] Any code provided is minimal, complete, and reproducible ([more on MREs](https://docs.github.com/en/get-started/writing-on-github/working-with-advanced-formatting/creating-and-highlighting-code-blocks))
- [x] Any code provided is properly formatted in code blocks, (no screenshot, [more on code blocks](https://docs.github.com/en/get-started/writing-on-github/working-with-advanced-formatting/creating-and-highlighting-code-blocks))
- [x] Any traceback provided is complete

Contributor guide

Open the contributing guide

Research direction

Start with GRPOTrainer._tokenize_prompts and _generate_single_turn, then run the SmolVLM diagnostic to compare processor-expanded and tokenizer-only prompt IDs. Verify that vLLM receives unexpanded IDs while the training forward pass retains expanded IDs, and confirm the rollout behavior against the attached CSV comparison.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
ai, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.