add all2all

Signed-off-by: angazenn <zengyanjia@huawei.com>

Author: angazenn

vllm_ascend/distributed/communicator.py

@@ -14,9 +14,10 @@
 # limitations under the License.
 # This file is a part of the vllm-ascend project.
 #
-from typing import Optional
+from typing import Optional, List
 
 import torch
+import torch.distributed as dist
 from torch.distributed import ProcessGroup
 from vllm.distributed.device_communicators.base_device_communicator import \
     DeviceCommunicatorBase
@@ -33,3 +34,58 @@ def __init__(self,
         # TODO(hz): Refer to CudaCommunicator's implementation to integrate PyHcclCommunicator
         # init device according to rank
         self.device = torch.npu.current_device()
+
+    def all_to_all(self,
+                   input_: torch.Tensor,
+                   scatter_dim: int = 0,
+                   gather_dim: int = -1,
+                   scatter_sizes: List[int] = None,
+                   gather_sizes: List[int] = None) -> torch.Tensor:
+
+        if scatter_dim < 0:
+            scatter_dim += input_.dim()
+        if gather_dim < 0:
+            gather_dim += input_.dim()
+
+        if scatter_sizes is not None and gather_sizes is not None:
+            input_list = [t.contiguous() for t in torch.split(input_, scatter_sizes, scatter_dim)]
+            output_list = []
+            tensor_shape_base = input_list[self.rank].size()
+            for i in range(self.world_size):
+                tensor_shape = list(tensor_shape_base)
+                tensor_shape[gather_dim] = gather_sizes[i]
+                output_list.append(torch.empty(tensor_shape, dtype=input_.dtype, device=input_.device))
+
+        else:
+            input_list = [
+                t.contiguous()
+                for t in torch.tensor_split(input_, self.world_size, scatter_dim)
+            ]
+            output_list = [torch.empty_like(input_list[i]) for i in range(self.world_size)]
+
+        dist.all_to_all(output_list, input_list, group=self.device_group)
+        output_tensor = torch.cat(output_list, dim=gather_dim).contiguous()
+        return output_tensor
+
+    def reduce_scatter(self,
+                       input_: torch.Tensor,
+                       scatter_dim: int = 0) -> torch.Tensor:
+
+        if scatter_dim < 0:
+            scatter_dim += input_.dim()
+        if scatter_dim != 0:
+            input_ = torch.transpose(input_, 0, scatter_dim)
+        dim_size = list(input_.size())
+        dim_size[0] = dim_size[0] // self.world_size
+        output_tensor = torch.empty(dim_size,
+                                    dtype=input_.dtype,
+                                    device=input_.device)
+        dist.reduce_scatter_tensor(output_tensor,
+                                   input_.contiguous(),
+                                   group=self.device_group)
+        if scatter_dim != 0:
+            output_tensor = torch.transpose(output_tensor, 0,
+                                            scatter_dim).contiguous()
+        return output_tensor
+
+

vllm_ascend/ops/fused_moe.py

@@ -627,6 +627,7 @@ def forward(self,
             if int(os.environ.get("VLLM_ENABLE_MC2", '0')  # type: ignore
                    ) == 1 and not is_prefill:
                 ...
+        else:
             elif int(os.environ.get("USING_LCCL_COM",
                                     '0')) == 1:  # type: ignore
                 hidden_states = get_dp_group().all_gather(
@@ -652,18 +653,19 @@ def forward(self,
             custom_routing_function=self.custom_routing_function,
             scoring_func=self.scoring_func,
             e_score_correction_bias=self.e_score_correction_bias,
-            is_prefill=is_prefill)
+            is_prefill=is_prefill,
+            dp_size=self.dp_size)
 
         if self.dp_size > 1:
             if int(os.environ.get("VLLM_ENABLE_MC2", '0')  # type: ignore
                    ) == 1 and not is_prefill:
                 ...
-            else:
-                final_hidden_states = dist._functional_collectives.reduce_scatter_tensor(
-                    final_hidden_states,
-                    "sum",
-                    scatter_dim=0,
-                    group=get_dp_group().device_group)
+        else:
+            final_hidden_states = dist._functional_collectives.reduce_scatter_tensor(
+                final_hidden_states,
+                "sum",
+                scatter_dim=0,
+                group=get_dp_group().device_group)
 
         if self.reduce_results and (self.tp_size > 1 or self.ep_size > 1):
             final_hidden_states = tensor_model_parallel_all_reduce(

vllm_ascend/patch/platform/patch_common/patch_distributed.py

@@ -16,16 +16,18 @@
 # limitations under the License.
 # Adapted from vllm/model_executor/models/qwen2_vl.py
 # This file is a part of the vllm-ascend project.
+from typing import Optional, List
 
 import torch
-import vllm
-import vllm.distributed
 from torch.distributed import ProcessGroup
 from torch.distributed.distributed_c10d import (Backend, PrefixStore,
                                                 _get_default_timeout,
                                                 is_nccl_available)
 from torch.distributed.rendezvous import rendezvous
+import vllm
 from vllm.config import ParallelConfig
+import vllm.distributed
+from vllm.distributed.parallel_state import GroupCoordinator
 
 
 def ascend_destroy_model_parallel():
@@ -176,6 +178,40 @@ def ascend_stateless_init_dp_group(self) -> "ProcessGroup":
     return dp_group
 
 
+class GroupCoordinatorPatch(GroupCoordinator):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def all_to_all(self,
+                   input_: torch.Tensor,
+                   scatter_dim: int = 0,
+                   gather_dim: int = -1,
+                   scatter_sizes: List[int] = None,
+                   gather_sizes: List[int] = None) -> torch.Tensor:
+        if self.world_size == 1:
+            return input_
+        assert -input_.dim() <= scatter_dim < input_.dim(), (
+            f"Invalid scatter dim ({scatter_dim}) for input tensor with shape {input_.size()}"
+        )
+        assert -input_.dim() <= gather_dim < input_.dim(), (
+            f"Invalid gather dim ({gather_dim}) for input tensor with shape {input_.size()}"
+        )
+        return self.device_communicator.all_to_all(input_, scatter_dim, gather_dim,
+                                                   scatter_sizes, gather_sizes)
+
+    def reduce_scatter(self,
+                       input_: torch.Tensor,
+                       scatter_dim: int = 0) -> torch.Tensor:
+        if self.world_size == 1:
+            return input_
+        assert -input_.dim() <= scatter_dim < input_.dim(), (
+            f"Invalid scatter dim ({scatter_dim}) for input tensor with shape {input_.size()}"
+        )
+        return self.device_communicator.reduce_scatter(input_, scatter_dim)
+
+
+vllm.distributed.parallel_state.GroupCoordinator = GroupCoordinatorPatch # Note: check the GroupCoordinator with online serving
 vllm.distributed.parallel_state.destroy_model_parallel = ascend_destroy_model_parallel
 vllm.distributed.stateless_init_torch_distributed_process_group = ascend_stateless_init_torch_distributed_process_group
 ParallelConfig.get_next_dp_init_port = parallel_config_get_dp_port

vllm_ascend/quantization/quant_config.py

@@ -317,14 +317,16 @@ def apply(
         scoring_func: str = "softmax",
         e_score_correction_bias: Optional[torch.Tensor] = None,
         is_prefill: bool = True,
+        dp_size: int = 1
         **kwargs,
     ) -> torch.Tensor:
         return self.quant_method.apply(layer, x, router_logits, top_k,
                                        renormalize, use_grouped_topk,
                                        global_num_experts, expert_map,
                                        topk_group, num_expert_group,
                                        custom_routing_function, scoring_func,
-                                       e_score_correction_bias, is_prefill)
+                                       e_score_correction_bias, is_prefill,
+                                       dp_size)
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         if hasattr(self.quant_method, "process_weights_after_loading"):

vllm_ascend/quantization/w8a8_dynamic.py

@@ -20,7 +20,8 @@
 
 import torch
 import torch_npu
-
+import torch.distributed as dist
+from vllm.distributed import GroupCoordinator
 from vllm_ascend.distributed.parallel_state import get_ep_group
 from vllm_ascend.ops.fused_moe import select_experts
 
@@ -201,6 +202,120 @@ def fused_experts_with_mc2(
     return hidden_states
 
 
+def fused_experts_with_all2all(hidden_states: torch.Tensor,
+                               w1: torch.Tensor,
+                               w1_scale: torch.Tensor,
+                               w2: torch.Tensor,
+                               w2_scale: torch.Tensor,
+                               topk_weights: torch.Tensor,
+                               topk_ids: torch.Tensor,
+                               top_k: int,
+                               expert_map: torch.Tensor = None,
+                               ep_group: GroupCoordinator = None,
+                               ):
+    original_shape = hidden_states.shape
+    if len(original_shape) == 3:
+        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
+
+    num_tokens, _ = hidden_states.shape
+    num_experts = w1.shape[0]
+    dtype = hidden_states.dtype
+    device = hidden_states.device
+
+    if expert_map is not None:
+        global_num_experts = len(expert_map)
+        local_num_experts = global_num_experts // ep_group.world_size
+        row_idx_len = num_tokens * top_k
+        row_idx = (torch.arange(0,
+                                row_idx_len,
+                                dtype=torch.int32,
+                                device=device).view(top_k, -1).permute(
+            1, 0).contiguous())
+        hidden_states, expanded_row_idx, expanded_expert_idx = torch_npu.npu_moe_init_routing(
+            hidden_states,
+            row_idx=row_idx,
+            expert_idx=topk_ids,
+            active_num=num_tokens)
+
+        global_expert_tokens = torch.bincount(expanded_expert_idx, minlength=global_num_experts)
+        scatter_sizes = global_expert_tokens.view(ep_group.world_size, -1).sum(-1)
+
+        gather_sizes = torch.empty_like(scatter_sizes)
+        dist.all_to_all_single(gather_sizes, scatter_sizes, group=ep_group.device_group)
+        scatter_size_list = scatter_sizes.cpu().tolist()
+        gather_size_list = gather_sizes.cpu().tolist()
+
+        expanded_expert_idx = expanded_expert_idx % local_num_experts
+        hidden_states = ep_group.all_to_all(hidden_states, 0, 0, scatter_size_list, gather_size_list)
+        local_expert_idx = ep_group.all_to_all(expanded_expert_idx, 0, 0, scatter_size_list, gather_size_list)
+
+        sorted_local_expert_idx, sorted_idx = torch.sort(local_expert_idx)
+
+        expert_tokens = torch_npu.npu_moe_compute_expert_tokens(sorted_local_expert_idx, local_num_experts).to(
+            torch.int64)
+
+        hidden_states = hidden_states[sorted_idx]
+        group_list_type = 0
+    else:
+        row_idx_len = num_tokens * top_k
+        row_idx = torch.arange(0,
+                               row_idx_len,
+                               dtype=torch.int32,
+                               device=topk_weights.device).view(
+                                   top_k, -1).permute(1, 0).contiguous()
+        hidden_states, expanded_row_idx, expanded_expert_idx = torch_npu.npu_moe_init_routing(
+            hidden_states,
+            row_idx=row_idx,
+            expert_idx=topk_ids,
+            active_num=num_tokens)
+
+        expert_tokens = torch_npu.npu_moe_compute_expert_tokens(
+            expanded_expert_idx, num_experts)
+        expert_tokens = expert_tokens.to(torch.int64)
+        group_list_type = 0
+
+    hidden_states_wrapper = [hidden_states]
+    del hidden_states
+
+    hidden_states = apply_mlp(hidden_states_wrapper,
+                              w1,
+                              w1_scale,
+                              w2,
+                              w2_scale,
+                              expert_tokens,
+                              group_list_type=group_list_type)
+
+    if expert_map is not None:
+        resorted_idx = torch.argsort(sorted_idx)
+        hidden_states = hidden_states[resorted_idx]
+        hidden_states = ep_group.all_to_all(hidden_states, 0, 0, gather_size_list, scatter_size_list)
+
+        final_hidden_states = torch_npu.npu_moe_finalize_routing(
+            hidden_states,
+            skip1=None,
+            skip2=None,
+            bias=None,
+            scales=topk_weights,
+            expanded_src_to_dst_row=expanded_row_idx,
+            export_for_source_row=topk_ids,
+        )
+    else:
+        # TODO: Reorder device memory 2 times here, replace the current
+        # implementation here when suitable operators become available.
+        final_hidden_states = torch_npu.npu_moe_finalize_routing(
+            hidden_states,
+            skip1=None,
+            skip2=None,
+            bias=None,
+            scales=topk_weights,
+            expanded_src_to_dst_row=expanded_row_idx,
+            export_for_source_row=topk_ids,
+        )
+    if len(original_shape) == 3:
+        final_hidden_states = final_hidden_states.view(original_shape)
+    return final_hidden_states
+
+
 def fused_experts(hidden_states: torch.Tensor,
                   w1: torch.Tensor,
                   w1_scale: torch.Tensor,
@@ -387,10 +502,10 @@ class AscendW8A8DynamicFusedMoEMethod:
     def __init__(self):
         self.transpose_weight = True
 
-        ep_group = get_ep_group()
+        self.ep_group = get_ep_group()
 
         try:
-            device_group = ep_group.device_group
+            device_group = self.ep_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"))
@@ -457,6 +572,7 @@ def apply(
         scoring_func: str = "softmax",
         e_score_correction_bias: Optional[torch.Tensor] = None,
         is_prefill: bool = True,
+        dp_size: int = 1,
         **kwargs,
     ) -> torch.Tensor:
         assert router_logits.shape[
@@ -503,7 +619,7 @@ def apply(
                 top_k=top_k,
                 expert_map=expert_map,
                 moe_all_to_all_group_name=self.moe_all_to_all_group_name)
-        else:
+        elif dp_size == 1:
             return fused_experts(hidden_states=x,
                                  w1=layer.w13_weight,
                                  w1_scale=layer.w13_weight_scale,
@@ -513,6 +629,17 @@ def apply(
                                  topk_ids=topk_ids,
                                  top_k=top_k,
                                  expert_map=expert_map)
+        else:
+            return fused_experts_with_all2all(hidden_states=x,
+                                              w1=layer.w13_weight,
+                                              w1_scale=layer.w13_weight_scale,
+                                              w2=layer.w2_weight,
+                                              w2_scale=layer.w2_weight_scale,
+                                              topk_weights=topk_weights,
+                                              topk_ids=topk_ids,
+                                              top_k=top_k,
+                                              expert_map=expert_map,
+                                              ep_group=self.ep_group)
 
     def process_weights_after_loading(self, layer):
         if self.transpose_weight: