[Bugfix] Fix external LB DP rank handling when replicas share nodes (#53743)

Signed-off-by: Tony Lin <tony.lin@intel.com>

docs/serving/data_parallel_deployment.md

@@ -126,6 +126,14 @@ vllm serve $MODEL --data-parallel-size 2 --data-parallel-rank 1 \
                   --data-parallel-address 10.99.48.128 --data-parallel-rpc-port 13345
 ```
 
+When using `--data-parallel-external-lb` with `--nnodes`, the data parallel rank
+can be inferred from `--node-rank` only when `--nnodes` is evenly divisible by
+`--data-parallel-size`. If it is not evenly divisible, automatic rank inference
+is rejected and `--data-parallel-rank` must be set explicitly. When multiple
+external-LB ranks are co-located on a node, assign each process a disjoint
+device set with `CUDA_VISIBLE_DEVICES` or `--device-ids`; devices are not
+partitioned automatically between processes.
+
 The coordinator process also runs in this scenario, co-located with the DP rank 0 engine.
 
 <figure markdown="1">

tests/test_config.py

@@ -1518,6 +1518,41 @@ def test_reconfigure_for_independent_dp_rank_on_multinode_dense_model():
     assert parallel_config.world_size == 8
 
 
+@pytest.mark.parametrize("data_parallel_size", [4, 8])
+def test_nnodes_within_dp_when_replicas_outnumber_nodes(data_parallel_size):
+    """A replica that fits on one node spans one node, never zero.
+
+    External LB pins ``data_parallel_size_local`` to 1, so the ratio rounds
+    down to 0 as soon as there are more DP replicas than nodes. Anything
+    dividing by ``nnodes_within_dp`` then raises ZeroDivisionError.
+    """
+    parallel_config = ParallelConfig(
+        data_parallel_size=data_parallel_size,
+        data_parallel_size_local=1,
+        data_parallel_external_lb=True,
+        distributed_executor_backend="mp",
+        nnodes=2,
+        node_rank=1,
+    )
+
+    assert parallel_config.nnodes_within_dp == 1
+    assert parallel_config.node_rank_within_dp == 0
+    assert parallel_config.local_world_size == parallel_config.world_size
+
+
+@pytest.mark.parametrize("nnodes", [4, 9])
+def test_nnodes_within_dp_rejects_uneven_internal_lb(nnodes):
+    parallel_config = ParallelConfig(
+        data_parallel_size=8,
+        data_parallel_size_local=1,
+        distributed_executor_backend="mp",
+        nnodes=nnodes,
+    )
+
+    with pytest.raises(ValueError, match="Invalid data parallel configuration"):
+        _ = parallel_config.nnodes_within_dp
+
+
 def test_draft_model_enables_async_scheduling_by_default():
     parallel_config = ParallelConfig(distributed_executor_backend="uni")
     model_config = ModelConfig("Qwen/Qwen3-0.6B", max_model_len=2048)

tests/v1/engine/test_engine_args.py

@@ -2,9 +2,11 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from argparse import ArgumentError
+from unittest.mock import patch
 
 import pytest
 
+from vllm.config import ModelConfig
 from vllm.engine.arg_utils import EngineArgs
 from vllm.usage.usage_lib import UsageContext
 from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -121,3 +123,76 @@ def test_data_parallel_start_rank_zero_infers_hybrid_lb():
 
     assert vllm_config.parallel_config.data_parallel_hybrid_lb is True
     assert vllm_config.parallel_config.data_parallel_rank == 0
+
+
+def test_external_lb_preserves_explicit_rank_when_dp_exceeds_nodes():
+    engine_args = EngineArgs(
+        model="facebook/opt-125m",
+        data_parallel_size=4,
+        data_parallel_rank=3,
+        data_parallel_external_lb=True,
+        nnodes=2,
+        node_rank=1,
+    )
+
+    with patch.object(ModelConfig, "is_moe", new=property(lambda self: True)):
+        vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
+
+    assert vllm_config.parallel_config.data_parallel_rank == 3
+
+
+@pytest.mark.parametrize(
+    ("data_parallel_size", "nnodes", "tensor_parallel_size"),
+    [(4, 2, 1), (2, 3, 3)],
+)
+def test_external_lb_requires_explicit_rank_when_nodes_are_not_evenly_partitioned(
+    data_parallel_size, nnodes, tensor_parallel_size
+):
+    engine_args = EngineArgs(
+        model="facebook/opt-125m",
+        data_parallel_size=data_parallel_size,
+        data_parallel_external_lb=True,
+        tensor_parallel_size=tensor_parallel_size,
+        nnodes=nnodes,
+        node_rank=1,
+    )
+
+    with (
+        patch.object(ModelConfig, "is_moe", new=property(lambda self: True)),
+        pytest.raises(
+            ValueError,
+            match="Set a unique `--data-parallel-rank`",
+        ),
+    ):
+        engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
+
+
+def test_external_lb_infers_rank_when_dp_does_not_exceed_nodes():
+    engine_args = EngineArgs(
+        model="facebook/opt-125m",
+        data_parallel_size=2,
+        data_parallel_external_lb=True,
+        nnodes=2,
+        node_rank=1,
+    )
+
+    with patch.object(ModelConfig, "is_moe", new=property(lambda self: True)):
+        vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
+
+    assert vllm_config.parallel_config.data_parallel_rank == 1
+
+
+def test_external_lb_infers_rank_for_multinode_replicas():
+    engine_args = EngineArgs(
+        model="facebook/opt-125m",
+        data_parallel_size=2,
+        data_parallel_external_lb=True,
+        tensor_parallel_size=12,
+        nnodes=4,
+        node_rank=1,
+    )
+
+    with patch.object(ModelConfig, "is_moe", new=property(lambda self: True)):
+        vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
+
+    assert vllm_config.parallel_config.data_parallel_rank == 0

vllm/config/parallel.py

@@ -760,12 +760,30 @@ class ParallelConfig:
 
     @property
     def nnodes_within_dp(self) -> int:
+        """Number of nodes one DP replica spans.
+
+        External LB pins ``data_parallel_size_local`` to 1, so the ratio
+        rounds down to 0 once DP replicas outnumber nodes. A replica that
+        does not span nodes still occupies exactly one.
+        """
         if self.nnodes == 1:
             return 1
         data_parallel_node_size = (
             self.data_parallel_size // self.data_parallel_size_local
         )
-        return self.nnodes // data_parallel_node_size
+        nnodes_within_dp = self.nnodes // data_parallel_node_size
+        if self.data_parallel_external_lb:
+            return max(nnodes_within_dp, 1)
+        if self.nnodes % data_parallel_node_size != 0:
+            raise ValueError(
+                "Invalid data parallel configuration: "
+                f"nnodes ({self.nnodes}) must be divisible by the number of "
+                "data parallel node groups "
+                f"({data_parallel_node_size} = data_parallel_size "
+                f"{self.data_parallel_size} / data_parallel_size_local "
+                f"{self.data_parallel_size_local})"
+            )
+        return nnodes_within_dp
 
     @property
     def local_world_size(self) -> int:

vllm/engine/arg_utils.py

@@ -2232,12 +2232,22 @@ class EngineArgs:
                 self.node_rank * local_world_size
             ) // world_size_within_dp
             if self.data_parallel_size > 1 and self.data_parallel_external_lb:
-                self.data_parallel_rank = inferred_data_parallel_rank
-                logger.info(
-                    "Inferred data_parallel_rank %d from node_rank %d for external lb",
-                    self.data_parallel_rank,
-                    self.node_rank,
-                )
+                if self.data_parallel_rank is None:
+                    if self.nnodes % self.data_parallel_size != 0:
+                        raise ValueError(
+                            "Invalid data-parallel launch options: "
+                            "`--node-rank` cannot unambiguously identify external "
+                            "data-parallel ranks when `--nnodes` is not divisible "
+                            "by `--data-parallel-size`. Set a unique "
+                            "`--data-parallel-rank` for each external-LB process."
+                        )
+                    self.data_parallel_rank = inferred_data_parallel_rank
+                    logger.info(
+                        "Inferred data_parallel_rank %d from node_rank %d "
+                        "for external lb",
+                        self.data_parallel_rank,
+                        self.node_rank,
+                    )
             elif self.data_parallel_size_local is None:
                 # Infer data parallel size local for internal dplb:
                 self.data_parallel_size_local = max(