[BugFix] Fix data parallel initialization and device management in NPUWorker and DPEngineCoreProc

Co-authored-by: rjg-lyh <[email protected]>

Signed-off-by: Yizhou Liu <[email protected]>

Author: yiz-liu

vllm_ascend/patch/platform/patch_common/patch_distributed.py

@@ -25,7 +25,8 @@
                                                 _get_default_timeout,
                                                 is_nccl_available)
 from torch.distributed.rendezvous import rendezvous
-from vllm.config import ParallelConfig
+from vllm.config import ParallelConfig, VllmConfig
+from vllm.v1.engine.core import DPEngineCoreProc
 
 
 def ascend_destroy_model_parallel():
@@ -171,7 +172,7 @@ def parallel_config_get_dp_port(self) -> int:
     return port
 
 
-def ascend_stateless_init_dp_group(self) -> "ProcessGroup":
+def stateless_init_dp_group(self) -> "ProcessGroup":
     # TODO(Yizhou): Currently we have to set the backend to gloo
     # because in vllm.config.ParallelConfig.has_unfinished_dp the
     # device is set to cpu. We need to fix this in the future.
@@ -187,6 +188,21 @@ def ascend_stateless_init_dp_group(self) -> "ProcessGroup":
     return dp_group
 
 
+def _init_data_parallel(self, vllm_config: VllmConfig):
+    # Configure NPUs and stateless process group for data parallel.
+    dp_rank = vllm_config.parallel_config.data_parallel_rank
+    dp_size = vllm_config.parallel_config.data_parallel_size
+    local_dp_rank = vllm_config.parallel_config.data_parallel_rank_local
+
+    assert dp_size > 1
+    assert 0 <= local_dp_rank <= dp_rank < dp_size
+
+    self.local_dp_rank = local_dp_rank
+    self.dp_group = vllm_config.parallel_config.stateless_init_dp_group()
+    self.current_wave = 0
+
+
 vllm.distributed.parallel_state.destroy_model_parallel = ascend_destroy_model_parallel
+DPEngineCoreProc._init_data_parallel = _init_data_parallel
 ParallelConfig.get_next_dp_init_port = parallel_config_get_dp_port
-ParallelConfig.stateless_init_dp_group = ascend_stateless_init_dp_group
+ParallelConfig.stateless_init_dp_group = stateless_init_dp_group

vllm_ascend/worker/worker_v1.py

@@ -72,6 +72,13 @@ def __init__(
                          rank=rank,
                          distributed_init_method=distributed_init_method,
                          is_driver_worker=is_driver_worker)
+
+        # NOTE(Yizhou): Since we do not set ASCEND_RT_VISIBLE_DEVICES in
+        # vllm_ascend, we need to set the device id manually.
+        local_dp_rank = self.vllm_config.parallel_config.data_parallel_rank_local
+        world_size = self.vllm_config.parallel_config.world_size
+        self.local_rank_across_dp = local_dp_rank * world_size + self.local_rank
+
         # Try to import mindie_turbo to accelerate vLLM inference.
         try_register_lib(
             "mindie_turbo",
@@ -98,7 +105,7 @@ def wake_up(self, tags: Optional[list[str]] = None) -> None:
 
     def init_device(self):
         if self.device_config.device.type == "npu":
-            self.device = torch.device(f"npu:{self.local_rank}")
+            self.device = torch.device(f"npu:{self.local_rank_across_dp}")
             NPUPlatform.set_device(self.device)
             NPUPlatform.empty_cache()
             self.init_npu_memory = NPUPlatform.mem_get_info()[0]