add unit test and refine EPHybridCommunicateGroup

Author: pkuzyc

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

@@ -619,34 +619,50 @@ def __init__(
         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']
+        self._ep_degree = dim_dict.get('expert', 1)
+        self._moe_sharding_degree = dim_dict.get('moe_sharding', 1)
+        self._moe_pp_degree = dim_dict.get('pipe', 1)
+        self._dp_degree = dim_dict.get('data', 1)
+        self._mp_degree = dim_dict.get('model', 1)
+        self._pp_degree = dim_dict.get('pipe', 1)
+        self._sharding_degree = dim_dict.get('sharding', 1)
+        self._sep_degree = dim_dict.get('sep', 1)
 
         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)
+        assert (
+            "moe_sharding" in moe_hybrid_group_names
+            and "expert" in moe_hybrid_group_names
+        )
 
         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",
-        ]
+        if hybrid_group_names.index("pipe") > hybrid_group_names.index(
+            "moe_sharding"
+        ):
+            dense_group_names = [
+                "moe_sharding",
+                "pipe",
+                "dense_sharding",
+                "data",
+                "sep",
+                "model",
+            ]
+        else:
+            dense_group_names = [
+                "pipe",
+                "moe_sharding",
+                "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

test/collective/fleet/hybrid_parallel_communicate_group.py

@@ -16,13 +16,14 @@
 
 import paddle
 from paddle.distributed import fleet
+from paddle.distributed.fleet.base import topology as tp
 
 
 class TestNewGroupAPI:
     def __init__(self):
         paddle.distributed.init_parallel_env()
         topo = fleet.CommunicateTopology(
-            ["data", "model", "sharding", "pipe"], [2, 1, 1, 1]
+            ["data", "sep", "model", "sharding", "pipe"], [2, 1, 1, 1, 1]
         )
         self.hcg = fleet.HybridCommunicateGroup(topo)
 
@@ -101,6 +102,47 @@ def test_all(self):
         print("test barrier api ok")
 
 
+class TestHybridEPGroup:
+    def __init__(self):
+        paddle.distributed.init_parallel_env()
+        group_names = [
+            "moe_sharding",
+            "sharding",
+            "pipe",
+            "sep",
+            "data",
+            "expert",
+            "model",
+        ]
+        dims = [1, 1, 1, 1, 1, 2, 2]
+
+        self.hcg = tp.EPHybridCommunicateGroup(group_names, dims)
+
+    def test_all(self):
+        global_rank = paddle.distributed.get_rank()
+
+        dp_rank = self.hcg.get_data_parallel_rank()
+        assert dp_rank == 0
+        assert self.hcg.get_expert_parallel_world_size() == 2
+        assert self.hcg.get_moe_sharding_parallel_world_size() == 1
+        assert self.hcg.get_model_parallel_world_size() == 2
+        assert self.hcg.get_expert_parallel_rank() == global_rank
+        assert self.hcg.get_moe_sharding_parallel_rank() == 0
+        assert self.hcg.get_expert_parallel_group_src_rank() == 0
+        assert (
+            self.hcg.get_moe_sharding_parallel_group_src_rank() == global_rank
+        )
+
+        moe_sharding_group = self.hcg.get_moe_sharding_parallel_group()
+        ep_group = self.hcg.get_expert_parallel_group()
+        mp_group = self.hcg.get_model_parallel_group()
+        assert moe_sharding_group.ranks == [global_rank]
+        assert ep_group.ranks == [0, 1]
+        assert mp_group.ranks == [0, 1]
+
+
 if __name__ == "__main__":
     gpt = TestNewGroupAPI()
     gpt.test_all()
+    ep_test = TestHybridEPGroup()
+    ep_test.test_all()