[Bugfix] Fix external LB DP rank handling when replicas share nodes (#53743)
Signed-off-by: Tony Lin <tony.lin@intel.com>
Tony Lin committed
Sep 21, 2026 at 23:19 UTC
dc479a2d839d5c474b36ac521dc5a72294bcb186
5 files changed
+153
-7
docs/serving/data_parallel_deployment.md
+8
@@ -126,6 +126,14 @@ vllm serve $MODEL --data-parallel-size 2 --data-parallel-rank 1 \
126
--data-parallel-address 10.99.48.128 --data-parallel-rpc-port 13345
127
```
128
129
+When using `--data-parallel-external-lb` with `--nnodes`, the data parallel rank
130
+can be inferred from `--node-rank` only when `--nnodes` is evenly divisible by
131
+`--data-parallel-size`. If it is not evenly divisible, automatic rank inference
132
+is rejected and `--data-parallel-rank` must be set explicitly. When multiple
133
+external-LB ranks are co-located on a node, assign each process a disjoint
134
+device set with `CUDA_VISIBLE_DEVICES` or `--device-ids`; devices are not
135
+partitioned automatically between processes.
136
+
137
The coordinator process also runs in this scenario, co-located with the DP rank 0 engine.
138
139
<figure markdown="1">
tests/test_config.py
+35
@@ -1518,6 +1518,41 @@ def test_reconfigure_for_independent_dp_rank_on_multinode_dense_model():
1518
assert parallel_config.world_size == 8
1519
1520
1521
+@pytest.mark.parametrize("data_parallel_size", [4, 8])
1522
+def test_nnodes_within_dp_when_replicas_outnumber_nodes(data_parallel_size):
1523
+ """A replica that fits on one node spans one node, never zero.
1524
+
1525
+ External LB pins ``data_parallel_size_local`` to 1, so the ratio rounds
1526
+ down to 0 as soon as there are more DP replicas than nodes. Anything
1527
+ dividing by ``nnodes_within_dp`` then raises ZeroDivisionError.
1528
+ """
1529
+ parallel_config = ParallelConfig(
1530
+ data_parallel_size=data_parallel_size,
1531
+ data_parallel_size_local=1,
1532
+ data_parallel_external_lb=True,
1533
+ distributed_executor_backend="mp",
1534
+ nnodes=2,
1535
+ node_rank=1,
1536
+ )
1537
+
1538
+ assert parallel_config.nnodes_within_dp == 1
1539
+ assert parallel_config.node_rank_within_dp == 0
1540
+ assert parallel_config.local_world_size == parallel_config.world_size
1541
+
1542
+
1543
+@pytest.mark.parametrize("nnodes", [4, 9])
1544
+def test_nnodes_within_dp_rejects_uneven_internal_lb(nnodes):
1545
+ parallel_config = ParallelConfig(
1546
+ data_parallel_size=8,
1547
+ data_parallel_size_local=1,
1548
+ distributed_executor_backend="mp",
1549
+ nnodes=nnodes,
1550
+ )
1551
+
1552
+ with pytest.raises(ValueError, match="Invalid data parallel configuration"):
1553
+ _ = parallel_config.nnodes_within_dp
1554
+
1555
+
1556
def test_draft_model_enables_async_scheduling_by_default():
1557
parallel_config = ParallelConfig(distributed_executor_backend="uni")
1558
model_config = ModelConfig("Qwen/Qwen3-0.6B", max_model_len=2048)
tests/v1/engine/test_engine_args.py
+75
@@ -2,9 +2,11 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from argparse import ArgumentError
5
+from unittest.mock import patch
6
7
import pytest
8
9
+from vllm.config import ModelConfig
10
from vllm.engine.arg_utils import EngineArgs
11
from vllm.usage.usage_lib import UsageContext
12
from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -121,3 +123,76 @@ def test_data_parallel_start_rank_zero_infers_hybrid_lb():
123
124
assert vllm_config.parallel_config.data_parallel_hybrid_lb is True
125
assert vllm_config.parallel_config.data_parallel_rank == 0
126
+
127
+
128
+def test_external_lb_preserves_explicit_rank_when_dp_exceeds_nodes():
129
+ engine_args = EngineArgs(
130
+ model="facebook/opt-125m",
131
+ data_parallel_size=4,
132
+ data_parallel_rank=3,
133
+ data_parallel_external_lb=True,
134
+ nnodes=2,
135
+ node_rank=1,
136
+ )
137
+
138
+ with patch.object(ModelConfig, "is_moe", new=property(lambda self: True)):
139
+ vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
140
+
141
+ assert vllm_config.parallel_config.data_parallel_rank == 3
142
+
143
+
144
+@pytest.mark.parametrize(
145
+ ("data_parallel_size", "nnodes", "tensor_parallel_size"),
146
+ [(4, 2, 1), (2, 3, 3)],
147
+)
148
+def test_external_lb_requires_explicit_rank_when_nodes_are_not_evenly_partitioned(
149
+ data_parallel_size, nnodes, tensor_parallel_size
150
+):
151
+ engine_args = EngineArgs(
152
+ model="facebook/opt-125m",
153
+ data_parallel_size=data_parallel_size,
154
+ data_parallel_external_lb=True,
155
+ tensor_parallel_size=tensor_parallel_size,
156
+ nnodes=nnodes,
157
+ node_rank=1,
158
+ )
159
+
160
+ with (
161
+ patch.object(ModelConfig, "is_moe", new=property(lambda self: True)),
162
+ pytest.raises(
163
+ ValueError,
164
+ match="Set a unique `--data-parallel-rank`",
165
+ ),
166
+ ):
167
+ engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
168
+
169
+
170
+def test_external_lb_infers_rank_when_dp_does_not_exceed_nodes():
171
+ engine_args = EngineArgs(
172
+ model="facebook/opt-125m",
173
+ data_parallel_size=2,
174
+ data_parallel_external_lb=True,
175
+ nnodes=2,
176
+ node_rank=1,
177
+ )
178
+
179
+ with patch.object(ModelConfig, "is_moe", new=property(lambda self: True)):
180
+ vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
181
+
182
+ assert vllm_config.parallel_config.data_parallel_rank == 1
183
+
184
+
185
+def test_external_lb_infers_rank_for_multinode_replicas():
186
+ engine_args = EngineArgs(
187
+ model="facebook/opt-125m",
188
+ data_parallel_size=2,
189
+ data_parallel_external_lb=True,
190
+ tensor_parallel_size=12,
191
+ nnodes=4,
192
+ node_rank=1,
193
+ )
194
+
195
+ with patch.object(ModelConfig, "is_moe", new=property(lambda self: True)):
196
+ vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
197
+
198
+ assert vllm_config.parallel_config.data_parallel_rank == 0
vllm/config/parallel.py
+19
-1
@@ -760,12 +760,30 @@ class ParallelConfig:
760
761
@property
762
def nnodes_within_dp(self) -> int:
763
+ """Number of nodes one DP replica spans.
764
+
765
+ External LB pins ``data_parallel_size_local`` to 1, so the ratio
766
+ rounds down to 0 once DP replicas outnumber nodes. A replica that
767
+ does not span nodes still occupies exactly one.
768
+ """
769
if self.nnodes == 1:
770
return 1
771
data_parallel_node_size = (
772
self.data_parallel_size // self.data_parallel_size_local
773
)
768
- return self.nnodes // data_parallel_node_size
774
+ nnodes_within_dp = self.nnodes // data_parallel_node_size
775
+ if self.data_parallel_external_lb:
776
+ return max(nnodes_within_dp, 1)
777
+ if self.nnodes % data_parallel_node_size != 0:
778
+ raise ValueError(
779
+ "Invalid data parallel configuration: "
780
+ f"nnodes ({self.nnodes}) must be divisible by the number of "
781
+ "data parallel node groups "
782
+ f"({data_parallel_node_size} = data_parallel_size "
783
+ f"{self.data_parallel_size} / data_parallel_size_local "
784
+ f"{self.data_parallel_size_local})"
785
+ )
786
+ return nnodes_within_dp
787
788
@property
789
def local_world_size(self) -> int:
vllm/engine/arg_utils.py
+16
-6
@@ -2232,12 +2232,22 @@ class EngineArgs:
2232
self.node_rank * local_world_size
2233
) // world_size_within_dp
2234
if self.data_parallel_size > 1 and self.data_parallel_external_lb:
2235
- self.data_parallel_rank = inferred_data_parallel_rank
2236
- logger.info(
2237
- "Inferred data_parallel_rank %d from node_rank %d for external lb",
2238
- self.data_parallel_rank,
2239
- self.node_rank,
2240
- )
2235
+ if self.data_parallel_rank is None:
2236
+ if self.nnodes % self.data_parallel_size != 0:
2237
+ raise ValueError(
2238
+ "Invalid data-parallel launch options: "
2239
+ "`--node-rank` cannot unambiguously identify external "
2240
+ "data-parallel ranks when `--nnodes` is not divisible "
2241
+ "by `--data-parallel-size`. Set a unique "
2242
+ "`--data-parallel-rank` for each external-LB process."
2243
+ )
2244
+ self.data_parallel_rank = inferred_data_parallel_rank
2245
+ logger.info(
2246
+ "Inferred data_parallel_rank %d from node_rank %d "
2247
+ "for external lb",
2248
+ self.data_parallel_rank,
2249
+ self.node_rank,
2250
+ )
2251
elif self.data_parallel_size_local is None:
2252
# Infer data parallel size local for internal dplb:
2253
self.data_parallel_size_local = max(