Update vllm_rollout_spmd.py

gdw439 · web-flow · commit 29370e485931 · 2025-06-15T13:38:53.000+08:00
diff --git a/verl/workers/rollout/vllm_rollout_spmd.py b/verl/workers/rollout/vllm_rollout_spmd.py
@@ -167,7 +167,7 @@ def generate_sequences(self, prompts: DataProto) -> DataProto:
                 )
         else:
             vllm_inputs = [
-                {"prompt_token_ids": list(raw_prompt_ids)} for raw_prompt_ids in non_tensor_batch.pop("raw_prompt_ids")
+                {"prompt_token_ids": list(raw_prompt_ids)} for raw_prompt_ids in batch_raw_prompt_ids
             ]
 
         # users can customize different sampling_params at different run

Original file line number	Diff line number	Diff line change
`@@ -167,7 +167,7 @@ def generate_sequences(self, prompts: DataProto) -> DataProto:`
`167`	`167`	`)`
`168`	`168`	`else:`
`169`	`169`	`vllm_inputs = [`
`170`		`- {"prompt_token_ids": list(raw_prompt_ids)} for raw_prompt_ids in non_tensor_batch.pop("raw_prompt_ids")`
	`170`	`+ {"prompt_token_ids": list(raw_prompt_ids)} for raw_prompt_ids in batch_raw_prompt_ids`
`171`	`171`	`]`
`172`	`172`
`173`	`173`	`# users can customize different sampling_params at different run`