[0.9.1][bugfix] fix mc2 op GroupCoordinator bug (#1711)

### What this PR does / why we need it?


### Does this PR introduce _any_ user-facing change?


### How was this patch tested?

Signed-off-by: zzzzwwjj <1183291235@qq.com>

Author: zzzzwwjj

vllm_ascend/distributed/parallel_state.py

@@ -0,0 +1,49 @@
+from typing import Optional
+
+import torch
+from vllm.distributed.parallel_state import (GroupCoordinator, get_world_group,
+                                             init_model_parallel_group)
+
+# Currently, mc2 op need their own group coordinator.
+_MC2: Optional[GroupCoordinator] = None
+
+
+def get_mc2_group() -> GroupCoordinator:
+    assert _MC2 is not None, ("mc2 group is not initialized")
+    return _MC2
+
+
+def model_parallel_initialized():
+    return (_MC2 is not None)
+
+
+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)
+
+    _MC2 = init_model_parallel_group(group_ranks,
+                                     get_world_group().local_rank,
+                                     backend,
+                                     group_name="mc2")
+
+
+def destroy_ascend_model_parallel():
+    global _MC2
+    if _MC2:
+        _MC2.destroy()
+    _MC2 = None

vllm_ascend/ops/fused_moe.py

@@ -39,6 +39,7 @@
 import vllm_ascend.envs as envs_ascend
 from vllm_ascend.ascend_config import get_ascend_config
 from vllm_ascend.ascend_forward_context import FusedMoEState
+from vllm_ascend.distributed.parallel_state import get_mc2_group
 from vllm_ascend.ops.expert_load_balancer import ExpertLoadBalancer
 from vllm_ascend.utils import (AscendSocVersion, dispose_tensor,
                                get_ascend_soc_version, npu_stream_switch,
@@ -125,7 +126,7 @@ def fused_experts_with_mc2(
     mc2_mask: Optional[torch.Tensor] = None,
 ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
     quant_mode = 0
-    ep_group = get_ep_group()
+    ep_group = get_mc2_group()
     ep_rank_id = ep_group.rank_in_group
     ep_world_size = ep_group.world_size
     tp_world_size = get_tp_group().world_size
@@ -878,7 +879,7 @@ def __init__(self, moe: MoEConfig = None):
         self.torchair_graph_enabled = ascend_config.torchair_graph_config.enabled
 
         try:
-            device_group = get_ep_group().device_group
+            device_group = get_mc2_group().device_group
             # TODO: Try local_rank = ep_group.rank_in_group
             local_rank = torch.distributed.get_rank(group=device_group)
             backend = device_group._get_backend(torch.device("npu"))

vllm_ascend/patch/platform/patch_common/patch_distributed.py

@@ -17,13 +17,34 @@
 # Adapted from vllm/model_executor/models/qwen2_vl.py
 # This file is a part of the vllm-ascend project.
 
+import vllm
 import vllm.envs as envs
 from torch.distributed import ProcessGroup
 from vllm.config import ParallelConfig
 from vllm.distributed.utils import \
     stateless_init_torch_distributed_process_group
 
 
+def ascend_destroy_model_parallel():
+    """Set the groups to none and destroy them."""
+    from vllm.distributed.parallel_state import _DP, _EP, _PP, _TP
+    if _TP:
+        _TP.destroy()
+    _TP = None
+    if _PP:
+        _PP.destroy()
+    _PP = None
+    if _DP:
+        _DP.destroy()
+    _DP = None
+    if _EP:
+        _EP.destroy()
+    _EP = None
+    from vllm_ascend.distributed.parallel_state import \
+        destroy_ascend_model_parallel
+    destroy_ascend_model_parallel()
+
+
 def parallel_config_get_dp_port(self) -> int:
     """
     We might need to initialize process groups in multiple
@@ -57,5 +78,6 @@ def stateless_init_dp_group(self) -> "ProcessGroup":
     return dp_group
 
 
+vllm.distributed.parallel_state.destroy_model_parallel = ascend_destroy_model_parallel
 ParallelConfig.get_next_dp_init_port = parallel_config_get_dp_port
 ParallelConfig.stateless_init_dp_group = stateless_init_dp_group

vllm_ascend/quantization/w8a8_dynamic.py

@@ -26,6 +26,7 @@
 import vllm_ascend.envs as ascend_envs
 from vllm_ascend.ascend_config import get_ascend_config
 from vllm_ascend.ascend_forward_context import FusedMoEState
+from vllm_ascend.distributed.parallel_state import get_mc2_group
 from vllm_ascend.ops.fused_moe import select_experts
 from vllm_ascend.utils import (ACL_FORMAT_FRACTAL_NZ, AscendSocVersion,
                                dispose_tensor, get_ascend_soc_version,
@@ -223,7 +224,7 @@ def fused_experts_with_mc2(
     if log2phy:
         topk_ids = log2phy[topk_ids]
     quant_mode = 2
-    ep_group = get_ep_group()
+    ep_group = get_mc2_group()
     ep_rank_id = ep_group.rank_in_group
     ep_world_size = ep_group.world_size
 
@@ -763,7 +764,7 @@ def __init__(self):
         self.enable_weight_nz_layout = ascend_config.enable_weight_nz_layout
 
         try:
-            device_group = self.ep_group.device_group
+            device_group = get_mc2_group().device_group
             # TODO: Try local_rank = ep_group.rank_in_group
             local_rank = torch.distributed.get_rank(group=device_group)
             backend = device_group._get_backend(torch.device("npu"))

vllm_ascend/worker/worker.py

@@ -49,6 +49,7 @@
 
 from vllm_ascend.ascend_config import init_ascend_config
 from vllm_ascend.device_allocator.camem import CaMemAllocator
+from vllm_ascend.distributed.parallel_state import init_ascend_model_parallel
 from vllm_ascend.platform import NPUPlatform
 from vllm_ascend.utils import init_ascend_soc_version, try_register_lib
 from vllm_ascend.worker.model_runner import NPUModelRunner
@@ -545,6 +546,8 @@ 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)
         ensure_kv_transfer_initialized(vllm_config)
 
 

vllm_ascend/worker/worker_v1.py

@@ -40,6 +40,7 @@
 
 from vllm_ascend.ascend_config import init_ascend_config
 from vllm_ascend.device_allocator.camem import CaMemAllocator
+from vllm_ascend.distributed.parallel_state import init_ascend_model_parallel
 from vllm_ascend.platform import NPUPlatform
 from vllm_ascend.utils import init_ascend_soc_version, try_register_lib
 from vllm_ascend.worker.model_runner_v1 import NPUModelRunner
@@ -260,6 +261,8 @@ 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)
         ensure_kv_transfer_initialized(self.vllm_config)
 
     def _init_profiler(self):