[vllm] fix: use data-parallel rank in vLLM ZMQ handles#6620
Conversation
There was a problem hiding this comment.
Code Review
This pull request introduces a helper function _resolve_vllm_weight_sync_local_rank to resolve the local rank for vLLM weight synchronization, taking into account tensor and data parallel configurations. This prevents socket collisions when multiple data-parallel workers run on the same node. The _get_zmq_handle method is updated to use this resolved rank, and comprehensive unit tests are added. Feedback suggests logging a warning when the function silently falls back to worker_local_rank if dp_local_rank cannot be resolved, as this could still lead to socket collisions.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if dp_local_rank is None: | ||
| return worker_local_rank |
There was a problem hiding this comment.
When dp_size > 1 or dp_local_size > 1, but dp_local_rank cannot be resolved (i.e., it is None), the function silently falls back to worker_local_rank. This will cause multiple DP workers on the same node to use the same socket, leading to silent socket collisions and weight synchronization failures.
Consider logging a warning when this fallback occurs to aid in debugging distributed configuration issues.
| if dp_local_rank is None: | |
| return worker_local_rank | |
| if dp_local_rank is None: | |
| logger.warning( | |
| "Data parallelism is enabled (dp_size=%d, dp_local_size=%d), but " | |
| "the local data-parallel rank could not be resolved from parallel_config. " | |
| "Falling back to worker_local_rank=%d, which may cause socket collisions.", | |
| dp_size, dp_local_size, worker_local_rank | |
| ) | |
| return worker_local_rank |
|
Could you also merge this into version 0.8.0? The same issue exists there too. |
Fixes #6615.
Summary
Why
The sender side builds the ZMQ handle from a rollout-local rank that spans the local TP/DP workers:
The receiver side was using the vLLM worker's
self.local_rankdirectly. In DP>1 colocated vLLM setups, that rank can be TP-local, so multiple DP workers on the same node can try to receive from the samerank-0/rank-1socket.This patch reconstructs the sender-side local rank as:
When DP is not enabled, the existing local-rank behavior is preserved.
Validation
python -m py_compile verl\workers\rollout\vllm_rollout\utils.py tests\workers\rollout\test_vllm_cli_args_on_cpu.pypython -m ruff check verl\workers\rollout\vllm_rollout\utils.py tests\workers\rollout\test_vllm_cli_args_on_cpu.pygit diff --check_resolve_vllm_weight_sync_local_rank()and_get_zmq_handle()I also attempted:
On my Windows machine this gets blocked by local vLLM packaging/native-extension availability after setting up a Python 3.12
.venvand installing the reachable Python dependencies. The same test should run in CI where vLLM is installed as a package.