[Dist][Bugfix] Fix mc2 process group

Signed-off-by: MengqingCao <[email protected]>

Author: MengqingCao

vllm_ascend/distributed/parallel_state.py

@@ -1,7 +1,7 @@
 from typing import Optional
 
 import torch
-from vllm.distributed.parallel_state import (GroupCoordinator, get_world_group,
+from vllm.distributed.parallel_state import (GroupCoordinator, get_world_group, get_ep_group,
                                              init_model_parallel_group)
 
 # Currently, mc2 op need their own group coordinator.
@@ -18,23 +18,16 @@ def model_parallel_initialized():
 
 
 def init_ascend_model_parallel(
-    expert_parallel_size: int = 1,
-    world_size: Optional[int] = None,
     backend: Optional[str] = None,
 ):
     if model_parallel_initialized():
         return
     assert torch.distributed.is_initialized()
-    world_size = world_size or torch.distributed.get_world_size()
     backend = backend or torch.distributed.get_backend(
         get_world_group().device_group)
-    num_expert_parallel_groups = world_size // expert_parallel_size
 
     global _MC2
-    group_ranks = []
-    for i in range(num_expert_parallel_groups):
-        ranks = list(range(i, world_size, num_expert_parallel_groups))
-        group_ranks.append(ranks)
+    group_ranks = get_ep_group().ranks
 
     _MC2 = init_model_parallel_group(group_ranks,
                                      get_world_group().local_rank,

vllm_ascend/worker/worker.py

@@ -546,8 +546,7 @@ def _init_worker_distributed_environment(
         ensure_model_parallel_initialized(
             parallel_config.tensor_parallel_size,
             parallel_config.pipeline_parallel_size)
-        init_ascend_model_parallel(parallel_config.expert_parallel_size,
-                                   parallel_config.world_size_across_dp)
+        init_ascend_model_parallel()
         ensure_kv_transfer_initialized(vllm_config)
 
 

vllm_ascend/worker/worker_v1.py

@@ -261,8 +261,7 @@ def _init_worker_distributed_environment(self) -> None:
         ensure_model_parallel_initialized(
             self.parallel_config.tensor_parallel_size,
             self.parallel_config.pipeline_parallel_size)
-        init_ascend_model_parallel(self.parallel_config.expert_parallel_size,
-                                   self.parallel_config.world_size_across_dp)
+        init_ascend_model_parallel()
         ensure_kv_transfer_initialized(self.vllm_config)
 
     def _init_profiler(self):