diff --git a/tests/v1/engine/test_engine_args.py b/tests/v1/engine/test_engine_args.py index c03c244bb1..5033f4768b 100644 --- a/tests/v1/engine/test_engine_args.py +++ b/tests/v1/engine/test_engine_args.py @@ -99,3 +99,20 @@ def test_mm_prefix_lm_raises_batched_tokens_floor(): vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER) assert vllm_config.scheduler_config.max_num_batched_tokens >= 2496 + + +def test_data_parallel_start_rank_zero_infers_hybrid_lb(): + """An explicit --data-parallel-start-rank 0 must be treated the same as + any other explicit start rank when inferring hybrid LB mode, not as + "unset" (regression test for a truthiness-vs-`is not None` bug). + """ + engine_args = EngineArgs( + model="facebook/opt-125m", + data_parallel_size=4, + data_parallel_size_local=2, + data_parallel_start_rank=0, + ) + vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER) + + assert vllm_config.parallel_config.data_parallel_hybrid_lb is True + assert vllm_config.parallel_config.data_parallel_rank == 0 diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 4e4448a09c..528c884d73 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -2158,7 +2158,7 @@ class EngineArgs: elif self.data_parallel_size_local is not None: data_parallel_size_local = self.data_parallel_size_local - if self.data_parallel_start_rank and not headless: + if self.data_parallel_start_rank is not None and not headless: # Infer hybrid LB mode. self.data_parallel_hybrid_lb = True @@ -2177,7 +2177,9 @@ class EngineArgs: self.data_parallel_hybrid_lb = False self.data_parallel_rank = ( - self.data_parallel_start_rank or inferred_data_parallel_rank + self.data_parallel_start_rank + if self.data_parallel_start_rank is not None + else inferred_data_parallel_rank ) if self.nnodes > 1: logger.info(