ITADN

GRPO VLM colocate: pre-tokenized prompt ids + raw image handoff mis-aligns with vLLM preprocessing (M-RoPE crash / silent prompt corruption)

#6401Closedguangboyu 创建于 2026-07-15
G
guangboyucommented
### 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 条评论