GRPO VLM colocate: pre-tokenized prompt ids + raw image handoff mis-aligns with vLLM preprocessing (M-RoPE crash / silent prompt corruption)
### System Info
- trl 1.8.0, transformers 5.13.0, vllm 0.23.0 (cu129), torch 2.11.0, peft 0.19.1
- Model: `Qwen/Qwen3-VL-2B-Instruct` (single image per prompt)
- GRPOTrainer, `use_vllm=True`, `vllm_mode="colocate"`, `vllm_enable_sleep_mode=True`
### Bug description
In colocate mode, `VLLMGeneration` passes **pre-tokenized prompt ids** (built by the
HF processor with image pads already expanded) together with the **raw PIL image**
to `llm.generate` (`trl/generation/vllm_generation.py`, the
`row = {"prompt_token_ids": ids}; row["multi_modal_data"] = {"image": ...}` branch).
vLLM then runs its own multimodal preprocessing on that image. Whenever the two
pipelines disagree about the image token expansion, the ids and vLLM's features
mis-align. We observed two distinct failure modes on a VQA dataset with
heterogeneous image sizes (VQAv2+GQA mix):
1. **Hard crash** in the training forward:
`RuntimeError: shape mismatch: value tensor of shape [3, 413] cannot be broadcast
to indexing result of shape [3, 347]` in `Qwen3VLModel.get_rope_index`.
Reproduces with `per_device_train_batch_size=8, gradient_accumulation_steps=4`
and also with accumulation 1 at batch 16; only uniform-sized images avoid it.
2. **Silent prompt corruption** (worse): after forcing uniform image sizes, training
runs but completions ignore the prompt's format instructions, and
`sampling/importance_sampling_ratio/mean` collapses to ~1e-6 (rollout vs trainer
logps disagree by ~14 nats/token) while rewards flatline at 0.
Notably, a dataset with *identical* image sizes and near-identical prompt lengths
(CLEVR, 480x320) trains fine — which can hide the bug during sandbox validation.
### Reproduction sketch
Conversational prompts `[{"role": "user", "content": "<question>"}]` + a separate
`"image"` column (PIL, mixed sizes, e.g. VQAv2), Qwen3-VL-2B, GRPOConfig with
`use_vllm=True` (colocate). Crash occurs within the first ~3 optimizer steps.
### Workaround that fixes both symptoms
Routing generation through vLLM's own chat interface with a custom `rollout_func`
(so vLLM builds the prompt ids itself, and the trainer's forward re-processes the
same PIL deterministically to identical ids):
```python
def rollout(prompts, trainer):
conversations = [prepare_multimodal_messages_vllm(p) for p in prompts]
params = SamplingParams(temperature=trainer.args.temperature,
max_tokens=trainer.args.max_completion_length, logprobs=0)
outs = trainer.vllm_generation.llm.chat(conversations, params, use_tqdm=False)
return {"prompt_ids": [list(o.prompt_token_ids) for o in outs],
"completion_ids": [list(o.outputs[0].token_ids) for o in outs],
"logprobs": [[lps[t].logprob for t, lps in
zip(o.outputs[0].token_ids, o.outputs[0].logprobs)]
for o in outs]}
```
After this change: importance sampling ratio ~0.74, format rewards restored, no
crashes across 2000+ steps on the same data.
Happy to provide full configs/logs — the project where this surfaced is
https://github.com/guangboyu/vqa-rlvr.
关闭于 2026-07-16 2 条评论