add hybrid expert parallel communication group

Author: pkuzyc

paddle/fluid/framework/distributed_strategy.proto

@@ -119,6 +119,8 @@ message HybridConfig {
   optional DygraphShardingConfig sharding_configs = 8;
   optional bool enable_optimizer_timer = 9 [ default = false ];
   optional bool split_norm_comm = 10 [ default = false ];
+  optional int32 ep_degree = 11 [ default = 1 ];
+  optional int32 moe_sharding_degree = 12 [ default = 1 ];
 }
 
 message AMPConfig {

python/paddle/distributed/fleet/base/topology.py

@@ -14,6 +14,7 @@
 from __future__ import annotations
 
 import collections
+import math
 import os
 from functools import reduce
 from itertools import product
@@ -341,11 +342,13 @@ def _check_sep_exist(self) -> None:
         assert self._sep_degree > 1, "sep not exist"
 
     def _set_comm_group(
-        self, parallel_method: str = "data"
+        self, parallel_method: str = "data", topo: CommunicateTopology = None
     ) -> tuple[list[int], Group]:
         parallel_group = []
         parallel_comm_group = None
-        parallel_groups = self._topo.get_comm_list(parallel_method)
+        if topo is None:
+            topo = self._topo
+        parallel_groups = topo.get_comm_list(parallel_method)
 
         group_nccl_comm_init_option = (
             g_pipeline_nccl_comm_init_option
@@ -370,11 +373,13 @@ def _set_comm_group(
         return parallel_group, parallel_comm_group
 
     def _set_check_group(
-        self, parallel_method: str = "data"
+        self, parallel_method: str = "data", topo: CommunicateTopology = None
     ) -> tuple[list[int], Group]:
         parallel_group = []
         parallel_comm_group = None
-        parallel_size = self._topo.get_dim(parallel_method)
+        if topo is None:
+            topo = self._topo
+        parallel_size = topo.get_dim(parallel_method)
         for idx in range(parallel_size):
             parallel_groups = self._topo.get_axis_list(parallel_method, idx)
             comm_group = paddle.distributed.new_group(ranks=parallel_groups)
@@ -563,6 +568,9 @@ def get_pp_mp_parallel_group(self) -> Group:
         self._check_sep_exist()
         return self._pp_mp_comm_group
 
+    def get_moe_sharding_parallel_world_size(self) -> int:
+        return 0
+
     def create_fuse_group(
         self, fused_strategy_list: list[str]
     ) -> tuple[list[list[int]], list[Group]] | tuple[list[int], Group]:
@@ -593,6 +601,274 @@ def create_fuse_group(
             return parallel_group[0], parallel_comm_group[0]
 
 
+class EPHybridCommunicateGroup(HybridCommunicateGroup):
+    def __init__(
+        self,
+        hybrid_group_names: list[str] = [
+            "pipe",
+            "moe_sharding",
+            "expert",
+            "data",
+            "sharding",
+            "sep",
+            "model",
+        ],
+        dims: list[int] = [1, 1, 1, 1, 1, 1, 1],
+    ) -> None:
+        self.nranks = paddle.distributed.get_world_size()
+        self.global_rank = paddle.distributed.get_rank()
+
+        dim_dict = dict(zip(hybrid_group_names, dims))
+        self._ep_degree = dim_dict['expert']
+        self._moe_sharding_degree = dim_dict['moe_sharding']
+        self._moe_pp_degree = dim_dict['pipe']
+        self._dp_degree = dim_dict['data']
+        self._mp_degree = dim_dict['model']
+        self._pp_degree = dim_dict['pipe']
+        self._sharding_degree = dim_dict['sharding']
+        self._sep_degree = dim_dict['sep']
+
+        moe_hybrid_group_names = []
+        moe_dims = []
+        for name, dim in zip(hybrid_group_names, dims):
+            if name in ["pipe", "moe_sharding", "expert"]:
+                moe_hybrid_group_names.append(name)
+                moe_dims.append(dim)
+
+        self._moe_topo = CommunicateTopology(moe_hybrid_group_names, moe_dims)
+        dim_dict["dense_sharding"] = (
+            dim_dict["sharding"] // dim_dict["moe_sharding"]
+        )
+        dense_group_names = [
+            "moe_sharding",
+            "pipe",
+            "dense_sharding",
+            "data",
+            "sep",
+            "model",
+        ]
+        dense_dims = [dim_dict[name] for name in dense_group_names]
+        self._dense_topo = CommunicateTopology(dense_group_names, dense_dims)
+        self._moe_topo._parent_hcg = self
+        self._dense_topo._parent_hcg = self
+        self._topo = self._dense_topo
+
+        self._data_parallel_id = self._get_parallel_id(self._dense_topo, "data")
+        self._model_parallel_id = self._get_parallel_id(
+            self._dense_topo, "model"
+        )
+        self._sharding_parallel_id = self._get_sharding_parallel_id()
+        self._sep_parallel_id = self._get_parallel_id(self._dense_topo, "sep")
+        self.stage_id = self._get_parallel_id(self._moe_topo, "pipe")
+        self._expert_parallel_id = self._get_parallel_id(
+            self._moe_topo, "expert"
+        )
+        self._moe_sharding_parallel_id = self._get_parallel_id(
+            self._moe_topo, "moe_sharding"
+        )
+
+        assert (
+            self._moe_pp_degree == self._pp_degree
+        ), f"Mismatch moe_pp_degree:{self._moe_pp_degree}, pp_degree:{self._pp_degree}."
+        assert (
+            self._topo._world_size == self._moe_topo._world_size
+        ), f"Mismatch world_size:{self._topo._world_size}, moe_world_size:{self._moe_topo._world_size}."
+        assert (
+            self._sep_degree == 1 and self._dp_degree == 1
+        ), f"sep_degree {self._sep_degree} and dp_degree {self._dp_degree} must be 1 in MoE."
+
+        self._pp_group, self._pp_comm_group = self._set_comm_group(
+            "pipe", self._moe_topo
+        )
+        paddle.distributed.all_reduce(
+            paddle.zeros([1], dtype="int32"),
+            op=paddle.distributed.ReduceOp.SUM,
+            group=self._pp_comm_group,
+        )
+        env_name = "FLAGS_eager_communication_connection"
+        if paddle.get_flags(env_name)[env_name]:
+            if self._pp_comm_group.nranks > 1:
+                self._pp_comm_group.process_group.eager_connect_ring_exchange()
+
+        # create comm group for expert parallel
+        self._ep_group, self._ep_comm_group = self._set_comm_group(
+            "expert", self._moe_topo
+        )
+
+        # create comm group for sharding parallel in MoE layer
+        self._moe_sharding_group, self._moe_sharding_comm_group = (
+            self._set_comm_group("moe_sharding", self._moe_topo)
+        )
+
+        # create comm group for data parallel
+        self._dp_group, self._dp_comm_group = self._set_comm_group(
+            "data", self._dense_topo
+        )
+
+        # create comm group for sep parallel
+        self._sep_group, self._sep_comm_group = self._set_comm_group(
+            "sep", self._dense_topo
+        )
+
+        # create comm group for model parallel
+        self._mp_group, self._mp_comm_group = self._set_comm_group(
+            "model", self._dense_topo
+        )
+
+        # create comm group for sharding parallel
+        self._sharding_group, self._sharding_comm_group = (
+            self.build_sharding_group(self._dense_topo)
+        )
+
+        # create global group for check inf_nan / clip global norm
+        self._check_group, self._check_comm_group = self._set_check_group(
+            "data", self._dense_topo
+        )
+        self.sharding_group, self.sharding_check_comm_group = (
+            self.build_sharding_group(self._dense_topo)
+        )
+
+        # (
+        #     self.sharding_check_group,
+        #     self.sharding_check_comm_group,
+        # ) = self._set_check_group("sharding")
+
+        # create p2p group
+        self.is_first_stage = self.stage_id == 0
+        self.is_last_stage = self.stage_id == (self._pp_degree - 1)
+
+        # create p2p_groups
+        if self._pp_degree > 1:
+            if paddle.framework.core.is_compiled_with_nccl():
+                check_nccl_version_for_p2p()
+            self._set_p2p_prev_next()
+            if _use_four_directions:
+                self._set_four_directions_p2p_group()
+
+        debug_str = (
+            f"HybridParallelInfo: rank_id: {self.global_rank}, mp_degree: {self._mp_degree}, "
+            f"sharding_degree: {self._sharding_degree}, pp_degree: {self._pp_degree}, dp_degree: {self._dp_degree}, sep_degree: {self._sep_degree}, "
+            f"ep_degree: {self._ep_degree}, moe_sharding_degree: {self._moe_sharding_degree}"
+        )
+        debug_str += f", mp_group: {self._mp_group},  sharding_group: {self._sharding_group}, pp_group: {self._pp_group}, dp_group: {self._dp_group}, sep_group: {self._sep_group}, check/clip group: {self._check_group}, ep_group: {self._ep_group}, moe_sharding_group: {self._moe_sharding_group}."
+        logger.info(debug_str)
+
+        global _HYBRID_PARALLEL_GROUP
+        _HYBRID_PARALLEL_GROUP = self
+
+    def build_sharding_group(self, topo):
+        parallel_group = []
+        parallel_comm_group = None
+
+        parallel_groups = self.merge_inner_comm_list(
+            topo, "moe_sharding", "dense_sharding"
+        )
+
+        group_nccl_comm_init_option = 0
+
+        for group in parallel_groups:
+            comm_group = paddle.distributed.new_group(
+                ranks=group,
+                nccl_comm_init_option=group_nccl_comm_init_option,
+            )
+            if self.global_rank in group:
+                parallel_group = group
+                parallel_comm_group = comm_group
+
+        assert len(parallel_group) > 0
+        assert parallel_comm_group is not None
+
+        logger.info(
+            f"Total {len(parallel_groups)} sharding comm group(s) create successfully!"
+        )
+        return parallel_group, parallel_comm_group
+
+    def merge_inner_comm_list(self, topo, outer_name, inner_name):
+        """
+        merge all inner communication list whose rank-id are in
+        the same outer communication list. E.g.:
+          outer_comm_list: [[0, 4], [1, 5]]
+          inner_comm_list: [[0, 2], [1, 3], [4, 6], [5, 7]]
+          => merged_inner_comm_list: [[0, 2, 4, 6], [1, 3, 5, 7]]
+        """
+        inner_axis = topo._parallel_names.index(inner_name)
+        outer_axis = topo._parallel_names.index(outer_name)
+        inner_comm_list = topo.get_comm_list(inner_name)
+
+        num_merged_groups = len(inner_comm_list) // topo._dims[outer_axis]
+        interval = (
+            math.prod(topo._dims[(outer_axis + 1) :]) // topo._dims[inner_axis]
+        )
+        assert num_merged_groups > 0 and interval > 0
+
+        merged_comm_list = []
+        for i in range(num_merged_groups):
+            comm = []
+            for j in range(topo._dims[outer_axis]):
+                assert i + j * interval < len(
+                    inner_comm_list
+                ), f"Unexpected error in merge_inner_comm_list, {i}, {j}, {interval}, {len(inner_comm_list)}"
+                comm += inner_comm_list[i + j * interval]
+            merged_comm_list.append(comm)
+
+        return merged_comm_list
+
+    def find_col_idx(self, comm_list, global_rank):
+        rows = len(comm_list)
+        cols = len(comm_list[0])
+        r = rows - 1
+        c = 0
+
+        while r >= 0 and c < cols:
+            current = comm_list[r][c]
+            if current == global_rank:
+                return c
+            elif current < global_rank:
+                c += 1
+            else:
+                r -= 1
+
+        return None
+
+    def _get_parallel_id(self, topo, parallel_type):
+        comm_list = topo.get_comm_list(parallel_type)
+        parallel_id = self.find_col_idx(comm_list, self.global_rank)
+        assert parallel_id is not None
+        return parallel_id
+
+    def _get_sharding_parallel_id(self):
+        sharding_comm_list = self.merge_inner_comm_list(
+            self._dense_topo, "moe_sharding", "dense_sharding"
+        )
+        parallel_id = self.find_col_idx(sharding_comm_list, self.global_rank)
+        assert parallel_id is not None
+        return parallel_id
+
+    def get_expert_parallel_rank(self) -> int:
+        return self._expert_parallel_id
+
+    def get_expert_parallel_world_size(self) -> int:
+        return self._ep_degree
+
+    def get_expert_parallel_group(self) -> Group:
+        return self._ep_comm_group
+
+    def get_expert_parallel_group_src_rank(self) -> int:
+        return self._ep_comm_group.ranks[0]
+
+    def get_moe_sharding_parallel_rank(self) -> int:
+        return self._moe_sharding_parallel_id
+
+    def get_moe_sharding_parallel_world_size(self) -> int:
+        return self._moe_sharding_degree
+
+    def get_moe_sharding_parallel_group(self) -> Group:
+        return self._moe_sharding_comm_group
+
+    def get_moe_sharding_parallel_group_src_rank(self) -> int:
+        return self._moe_sharding_comm_group.ranks[0]
+
+
 class _CommunicateGroup:
     """tmp for static"""