[LLM] add moe parallel groups

paddlenlp/trainer/training_args.py

@@ -670,6 +670,14 @@
             )
         },
     )
+    expert_parallel_degree: int = field(
+        default=-1,
+        metadata={"help": ("The paddle expert data parallel strategy.")},
+    )
+    expert_tensor_parallel_degree: int = field(
+        default=-1,
+        metadata={"help": ("The paddle expert tensor parallel strategy. Currently is not supported. DO NOT SET.")},
+    )
     data_parallel_config: str = field(
         default="",
         metadata={
@@ -1123,6 +1131,13 @@
             sep_parallel_degree = max(self.sep_parallel_degree, 1)
             context_parallel_degree = max(self.context_parallel_degree, 1)
             pipeline_parallel_degree = max(self.pipeline_parallel_degree, 1)
+            expert_parallel_degree = max(self.expert_parallel_degree, 1)
+            expert_tensor_parallel_degree = max(self.expert_tensor_parallel_degree, 1)
+
+            # TODO(@gexiao): support expert_tensor_parallel_degree > 1 in the future
+            assert (
+                expert_tensor_parallel_degree == 1
+            ), f"Currently only support expert_tensor_parallel_degree=1, but got expert_tensor_parallel_degree of {expert_tensor_parallel_degree}"
 
             assert (
                 world_size % (self.tensor_parallel_degree * self.pipeline_parallel_degree) == 0
@@ -1146,6 +1161,11 @@
                 logger.warning("sharding_parallel_degree=1 means no sharding, please set sharding to empty!")
                 self.sharding = []
 
+            if sharding_parallel_degree > 1:
+                assert (
+                    sharding_parallel_degree % expert_parallel_degree == 0
+                ), f"sharding_parallel_degree should be divided by expert_parallel_degree, current sharding_parallel_degree: {sharding_parallel_degree}, expert_parallel_degree: {expert_parallel_degree}."
+
             self.data_parallel_degree = world_size // (
                 sharding_parallel_degree
                 * tensor_parallel_degree
@@ -1154,19 +1174,27 @@
                 * pipeline_parallel_degree
             )
 
+            assert not (
+                self.data_parallel_degree > 1 and expert_parallel_degree > 1
+            ), f"Currently only support use expert_data_parallel strategy together with sharding_parallel strategy, but not with data_parallel strategy. Currently data_parallel_degree is {self.data_parallel_degree}."
+
             if (
                 sharding_parallel_degree > 1
                 or tensor_parallel_degree > 1
                 or pipeline_parallel_degree > 1
                 or self.sep_parallel_degree > 1
                 or self.context_parallel_degree > 1
+                or expert_parallel_degree > 1
+                or expert_tensor_parallel_degree > 1
             ):
                 self.use_hybrid_parallel = True
                 self.sharding_parallel_degree = sharding_parallel_degree
                 self.tensor_parallel_degree = tensor_parallel_degree
                 self.pipeline_parallel_degree = pipeline_parallel_degree
                 self.sep_parallel_degree = sep_parallel_degree
                 self.context_parallel_degree = context_parallel_degree
+                self.expert_parallel_degree = expert_parallel_degree
+                self.expert_tensor_parallel_degree = expert_tensor_parallel_degree
 
             if not self.use_hybrid_parallel:
                 self.sharding = []
@@ -1175,6 +1203,8 @@
                 self.pipeline_parallel_degree = -1
                 self.sep_parallel_degree = -1
                 self.context_parallel_degree = -1
+                self.expert_parallel_degree = -1
+                self.expert_tensor_parallel_degree = -1
 
         if self.hybrid_parallel_topo_order is None:
             self.hybrid_parallel_topo_order = "pp_first"
@@ -1530,6 +1560,9 @@
                 fleet.init(is_collective=True, strategy=strategy)
                 logger.info(strategy)
 
+                if self.expert_parallel_degree > 1:
+                    self.add_moe_comm_group()
+
         elif self.enable_auto_parallel:
             self.tensor_parallel_degree = max(self.tensor_parallel_degree, 1)
             self.sep_parallel_degree = max(self.sep_parallel_degree, 1)
@@ -1738,6 +1771,10 @@
                 order = ["pp", "dp", "sharding", "sep", "mp"]
             elif self.hybrid_parallel_topo_order == "sharding_first":
                 order = ["dp", "sharding", "pp", "sep", "mp"]
+                if self.expert_parallel_degree > 1:
+                    logger.warning(
+                        "Currently using sharding_first topo order, but pp_first is recommended when using experts parallel for performance."
+                    )
 
             strategy = fleet.DistributedStrategy()
             strategy.hybrid_configs = {
@@ -1877,6 +1914,37 @@
                 )
                 self.pdc_download_ckpt = False
 
+    def add_moe_comm_group(self):
+        hcg = fleet.get_hybrid_communicate_group()
+        topo = hcg._topo
+        sharding_parallel_groups = topo.get_comm_list("sharding")
+        experts_replicas = self.sharding_parallel_degree // self.expert_parallel_degree
+
+        # init experts groups inside all sharding groups
+        for ranks_in_current_sharding_group in sharding_parallel_groups:
+            # init experts parallel groups (dispatch & combine)
+            for i in range(experts_replicas):
+                rank_indices = list(range(i * self.expert_parallel_degree, (i + 1) * self.expert_parallel_degree))
+                ranks = [ranks_in_current_sharding_group[i] for i in rank_indices]
+                group = dist.new_group(ranks=ranks)
+                if dist.get_rank() in ranks:
+                    assert not hasattr(hcg, "expert_parallel_group"), "expert_parallel_group can not be set repeate"
+                    setattr(hcg, "expert_parallel_group", group)
+
+            # init experts gradients comm groups
+            for i in range(self.expert_parallel_degree):
+                rank_indices = list(range(i, self.sharding_parallel_degree, self.expert_parallel_degree))
+                ranks = [ranks_in_current_sharding_group[i] for i in rank_indices]
+                group = dist.new_group(ranks=ranks)
+                if dist.get_rank() in ranks:
+                    assert not hasattr(hcg, "expert_grad_comm_group"), "expert_grad_comm_group can not be set repeate"
+                    setattr(hcg, "expert_grad_comm_group", group)
+
+        assert hasattr(hcg, "expert_parallel_group") and hasattr(hcg, "expert_grad_comm_group")
+        logger.info(
+            f"experts groups are created, expert_parallel_group: {hcg.expert_parallel_group}, expert_grad_comm_group: {hcg.expert_grad_comm_group}"
+        )
+
     def __str__(self):
         self_as_dict = asdict(self)
         self_as_dict = {k: f"<{k.upper()}>" if k.endswith("_token") else v for k, v in self_as_dict.items()}