[Bugfix] Account for PCP in multi-node world size validation (#55111)
Signed-off-by: DebugSy <532151398@qq.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Simon committed
Sep 3, 2026 at 18:52 UTC
ee0a4c46ae771af034ae676d44e932469586e7be
2 files changed
+25
-8
tests/engine/test_arg_utils.py
+19
@@ -464,6 +464,25 @@ def test_attention_config():
464
engine_args.create_engine_config()
465
466
467
+def test_multi_node_world_size_includes_pcp(monkeypatch):
468
+ """PCP expands the process world size, so the --nnodes divisibility check
469
+ must include it. Without this, TP=1/PCP=2 over 2 nodes computes a world
470
+ size of 1 and the launch is rejected before the engine starts."""
471
+ import vllm.config.vllm
472
+
473
+ # PCP requires the V2 model runner, which is gated on Triton.
474
+ monkeypatch.setattr(vllm.config.vllm, "HAS_TRITON", True)
475
+
476
+ engine_args = EngineArgs(
477
+ model="facebook/opt-125m",
478
+ tensor_parallel_size=1,
479
+ prefill_context_parallel_size=2,
480
+ nnodes=2,
481
+ )
482
+ vllm_config = engine_args.create_engine_config()
483
+ assert vllm_config.parallel_config.world_size == 2
484
+
485
+
486
def test_prefix_cache_default():
487
parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
488
args = parser.parse_args([])
vllm/engine/arg_utils.py
+6
-8
@@ -2121,21 +2121,19 @@ class EngineArgs:
2121
)
2122
inferred_data_parallel_rank = 0
2123
if self.nnodes > 1:
2124
- world_size = (
2125
- self.data_parallel_size
2126
- * self.pipeline_parallel_size
2127
- * self.tensor_parallel_size
2128
- )
2124
world_size_within_dp = (
2130
- self.pipeline_parallel_size * self.tensor_parallel_size
2125
+ self.pipeline_parallel_size
2126
+ * self.tensor_parallel_size
2127
+ * self.prefill_context_parallel_size
2128
)
2129
+ world_size = self.data_parallel_size * world_size_within_dp
2130
if world_size % self.nnodes != 0:
2131
raise ValueError(
2132
"Invalid data-parallel launch options: "
2133
f"`--nnodes {self.nnodes}` must evenly divide the total "
2134
f"world size ({world_size}). Adjust `--nnodes`, "
2137
- "`--data-parallel-size`, `--pipeline-parallel-size`, or "
2138
- "`--tensor-parallel-size`."
2135
+ "`--data-parallel-size`, `--pipeline-parallel-size`, "
2136
+ "`--tensor-parallel-size`, or `--prefill-context-parallel-size`."
2137
)
2138
if not 0 <= self.node_rank < self.nnodes:
2139
raise ValueError(