diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index b8d0206b7a..7338eefbbb 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -464,6 +464,25 @@ def test_attention_config(): engine_args.create_engine_config() +def test_multi_node_world_size_includes_pcp(monkeypatch): + """PCP expands the process world size, so the --nnodes divisibility check + must include it. Without this, TP=1/PCP=2 over 2 nodes computes a world + size of 1 and the launch is rejected before the engine starts.""" + import vllm.config.vllm + + # PCP requires the V2 model runner, which is gated on Triton. + monkeypatch.setattr(vllm.config.vllm, "HAS_TRITON", True) + + engine_args = EngineArgs( + model="facebook/opt-125m", + tensor_parallel_size=1, + prefill_context_parallel_size=2, + nnodes=2, + ) + vllm_config = engine_args.create_engine_config() + assert vllm_config.parallel_config.world_size == 2 + + def test_prefix_cache_default(): parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) args = parser.parse_args([]) diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index e89360505f..a14d76a3dd 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -2121,21 +2121,19 @@ class EngineArgs: ) inferred_data_parallel_rank = 0 if self.nnodes > 1: - world_size = ( - self.data_parallel_size - * self.pipeline_parallel_size - * self.tensor_parallel_size - ) world_size_within_dp = ( - self.pipeline_parallel_size * self.tensor_parallel_size + self.pipeline_parallel_size + * self.tensor_parallel_size + * self.prefill_context_parallel_size ) + world_size = self.data_parallel_size * world_size_within_dp if world_size % self.nnodes != 0: raise ValueError( "Invalid data-parallel launch options: " f"`--nnodes {self.nnodes}` must evenly divide the total " f"world size ({world_size}). Adjust `--nnodes`, " - "`--data-parallel-size`, `--pipeline-parallel-size`, or " - "`--tensor-parallel-size`." + "`--data-parallel-size`, `--pipeline-parallel-size`, " + "`--tensor-parallel-size`, or `--prefill-context-parallel-size`." ) if not 0 <= self.node_rank < self.nnodes: raise ValueError(