Support DP+TP+EP hybrid parallel deployment strategy

Author: carryyu

custom_ops/gpu_ops/moe/ep_moe_prefill_func.cu

@@ -43,11 +43,16 @@
       __VA_ARGS__                                                                       \
       break;                                                                            \
     }                                                                                   \
-    case 48: {                                                 \
-        constexpr size_t NUM_EXPERTS_PER_RANK = 48;            \
-        __VA_ARGS__                                            \
-        break;                                                 \
-    }   \
+    case 32: {                                                                          \
+      constexpr size_t NUM_EXPERTS_PER_RANK = 32;                                       \
+      __VA_ARGS__                                                                       \
+      break;                                                                            \
+    }                                                                                   \
+    case 48: {                                                                          \
+        constexpr size_t NUM_EXPERTS_PER_RANK = 48;                                     \
+        __VA_ARGS__                                                                     \
+        break;                                                                          \
+    }                                                                                   \
     case 64: {                                                                          \
       constexpr size_t NUM_EXPERTS_PER_RANK = 64;                                       \
       __VA_ARGS__                                                                       \

custom_ops/gpu_ops/save_with_output_msg.cc

@@ -105,7 +105,8 @@ void SaveOutMmsg(const paddle::Tensor& x,
                  int64_t rank_id,
                  int msg_queue_id,
                  bool save_each_rank) {
-    if (!save_each_rank && rank_id > 0) {
+    // don't use save_each_rank now!
+    if (rank_id > 0) {
         return;
     }
     if (x.place() == paddle::CPUPlace()) {

fastdeploy/config.py

@@ -22,6 +22,7 @@
 from enum import Enum
 from typing import Literal, Optional, Union
 
+import paddle.distributed as dist
 from paddleformers.transformers.configuration_utils import PretrainedConfig
 
 import fastdeploy
@@ -276,7 +277,10 @@ def __init__(
                 setattr(self, key, value)
 
         # currently, the expert parallel size is equal data parallel size
-        self.expert_parallel_size = self.data_parallel_size
+        if self.enable_expert_parallel:
+            self.expert_parallel_size = self.data_parallel_size * self.tensor_parallel_size
+        else:
+            self.expert_parallel_size = 1
         self.use_ep = self.expert_parallel_size > 1
         if self.splitwise_role == "mixed":
             self.moe_phase = MoEPhase(phase="prefill")
@@ -297,6 +301,22 @@ def __init__(
         else:
             self.pd_disaggregation_mode = "None"
 
+    def set_tp_group(self):
+        # different tp group id
+        # prevent different tp_groups using the same group_id
+        dist.collective._set_custom_gid(self.data_parallel_rank + 100)
+        self.tp_group = dist.new_group(
+            range(
+                self.data_parallel_rank * self.tensor_parallel_size,
+                (self.data_parallel_rank + 1) * self.tensor_parallel_size,
+            )
+        )
+        # same ep group id
+        dist.collective._set_custom_gid(self.data_parallel_size + 100)
+        logger.info(
+            f"data_parallel_size: {self.data_parallel_size}, tensor_parallel_size: {self.tensor_parallel_size}, expert_parallel_size: {self.expert_parallel_size}, data_parallel_rank: {self.data_parallel_rank}, tensor_parallel_rank: {self.tensor_parallel_rank}, expert_parallel_rank: {self.expert_parallel_rank}, tp_group: {self.tp_group}."
+        )
+
     def print(self):
         """
         print all config

fastdeploy/distributed/communication.py

@@ -47,15 +47,20 @@ def use_custom_allreduce(custom_all_reduce_max_bytes: int = 8192 * 1024):
     @paddle.jit.marker.unified
     def tensor_model_parallel_all_reduce(
         input_: paddle.Tensor,
+        group_: paddle.distributed.communication.group.Group = None,
     ) -> paddle.Tensor:
         """All-reduce the input tensor across model parallel group."""
         global _TP_AR
         if _TP_AR is not None and _TP_AR.should_custom_ar(input_):
+            # TODO: supports different_group custom allreduce
             _TP_AR.custom_all_reduce(input_)
         elif paddle.in_dynamic_mode():
-            hcg = fleet.get_hybrid_communicate_group()
-            mp_group = hcg.get_model_parallel_group()
-            dist.all_reduce(input_, group=mp_group)
+            if group_ is not None:
+                dist.all_reduce(input_, group=group_)
+            else:
+                hcg = fleet.get_hybrid_communicate_group()
+                mp_group = hcg.get_model_parallel_group()
+                dist.all_reduce(input_, group=mp_group)
         else:
             dist.all_reduce(input_)
 

fastdeploy/engine/config.py

@@ -192,12 +192,7 @@ def __init__(
         if self.enable_mm:
             self.max_prefill_batch = 1  # TODO:当前多模prefill阶段只支持并行度为1,待优化
 
-        # TODO(@wufeisheng): TP and EP need to be supported simultaneously.
-        assert (self.tensor_parallel_size == 1 and self.parallel_config.expert_parallel_size >= 1) or (
-            self.tensor_parallel_size >= 1 and self.parallel_config.expert_parallel_size == 1
-        ), "TP and EP cannot be enabled at the same time"
-
-        num_ranks = self.tensor_parallel_size * self.parallel_config.expert_parallel_size
+        num_ranks = self.tensor_parallel_size * self.parallel_config.data_parallel_size
         self.max_chips_per_node = 16 if current_platform.is_iluvatar() else 8
         if num_ranks > self.max_chips_per_node:
             self.worker_num_per_node = self.max_chips_per_node

fastdeploy/model_executor/layers/embeddings.py

@@ -57,43 +57,37 @@ def __init__(
         hcg = fleet.get_hybrid_communicate_group()
         self.mp_rank: int = hcg.get_model_parallel_rank()
         self.column_cut = False
-        self.world_size: int = hcg.get_model_parallel_world_size()
-        self.ring_id: int = hcg.get_model_parallel_group().id
-        self.use_ep: bool = fd_config.parallel_config.use_ep
+        self.world_size: int = fd_config.parallel_config.tensor_parallel_size
+        self.tensor_parallel_rank = fd_config.parallel_config.tensor_parallel_rank
+        self.tp_group = fd_config.parallel_config.tp_group
         self.hidden_dropout_prob: float = fd_config.model_config.hidden_dropout_prob
         self.initializer_range: float = fd_config.model_config.initializer_range
         self.max_position_embeddings: int = fd_config.model_config.max_position_embeddings
         self.tie_word_embeddings: bool = fd_config.model_config.tie_word_embeddings
         self.params_dtype: str = params_dtype
 
-        if self.use_ep:
-            self.embeddings = nn.Embedding(
+        if not self.column_cut:
+            self.embeddings = fleet.meta_parallel.VocabParallelEmbedding(
                 num_embeddings,
                 embedding_dim,
+                mp_group=self.tp_group,
+                weight_attr=paddle.ParamAttr(
+                    initializer=nn.initializer.Normal(mean=0.0, std=self.initializer_range),
+                ),
             )
+            if self.world_size > 1:
+                set_weight_attrs(self.embeddings.weight, {"output_dim": False})
         else:
-            if not self.column_cut:
-                self.embeddings = fleet.meta_parallel.VocabParallelEmbedding(
-                    num_embeddings,
-                    embedding_dim,
-                    mp_group=fleet.get_hybrid_communicate_group().get_model_parallel_group(),
-                    weight_attr=paddle.ParamAttr(
-                        initializer=nn.initializer.Normal(mean=0.0, std=self.initializer_range),
-                    ),
-                )
-                if self.world_size > 1:
-                    set_weight_attrs(self.embeddings.weight, {"output_dim": False})
-            else:
-                # column cut embedding
-                self.embeddings = nn.Embedding(
-                    num_embeddings,
-                    embedding_dim // self.world_size,
-                )
-
-                self.embeddings.weight.is_distributed = True
-                self.embeddings.weight.split_axis = 1
-                if self.world_size > 1:
-                    set_weight_attrs(self.embeddings.weight, {"output_dim": True})
+            # column cut embedding
+            self.embeddings = nn.Embedding(
+                num_embeddings,
+                embedding_dim // self.world_size,
+            )
+
+            self.embeddings.weight.is_distributed = True
+            self.embeddings.weight.split_axis = 1
+            if self.world_size > 1:
+                set_weight_attrs(self.embeddings.weight, {"output_dim": True})
 
         self.prefix = prefix
         self.dropout = nn.Dropout(self.hidden_dropout_prob)
@@ -125,20 +119,17 @@ def forward(self, ids_remove_padding=None) -> paddle.Tensor:
         Returns:
             Tensor: Embedded tensor representation of the input IDs.
         """
-        if self.use_ep:
+        if self.column_cut:
             input_embedings = self.embeddings(ids_remove_padding)
+            inputs_embeds_temp = []
+            paddle.distributed.all_gather(
+                inputs_embeds_temp,
+                input_embedings,
+                group=self.tp_group,
+                sync_op=True,
+            )
+            input_embedings = paddle.concat(inputs_embeds_temp, -1)
         else:
-            if self.column_cut:
-                input_embedings = self.embeddings(ids_remove_padding)
-                inputs_embeds_temp = []
-                paddle.distributed.all_gather(
-                    inputs_embeds_temp,
-                    input_embedings,
-                    group=fleet.get_hybrid_communicate_group().get_model_parallel_group(),
-                    sync_op=True,
-                )
-                input_embedings = paddle.concat(inputs_embeds_temp, -1)
-            else:
-                input_embedings = self.embeddings(ids_remove_padding)
+            input_embedings = self.embeddings(ids_remove_padding)
 
         return input_embedings

fastdeploy/model_executor/layers/linear.py

@@ -670,6 +670,7 @@ def __init__(
         self.fd_config = fd_config
         self.skip_quant = False
         self.nranks = fd_config.parallel_config.tensor_parallel_size
+        self.tp_group = fd_config.parallel_config.tp_group
         self.hidden_size = fd_config.model_config.hidden_size
         self.head_dim = fd_config.model_config.head_dim
         self.num_heads = fd_config.model_config.num_attention_heads // self.nranks
@@ -719,7 +720,7 @@ def forward_cuda(self, x: paddle.Tensor) -> paddle.Tensor:
             out = paddle.matmul(x, self.weight)
 
         if self.reduce_results and self.nranks > 1:
-            tensor_model_parallel_all_reduce(out)
+            tensor_model_parallel_all_reduce(out, self.tp_group)
 
         return out
 

fastdeploy/model_executor/layers/lm_head.py

@@ -58,7 +58,7 @@ def __init__(
             self.bias_key: Optional[str] = prefix + ".bias"
         else:
             self.bias_key: Optional[str] = None
-        self.use_ep: bool = fd_config.parallel_config.use_ep
+        self.tp_group = fd_config.parallel_config.tp_group
         self.column_cut = True
         self.nranks = fd_config.parallel_config.tensor_parallel_size
 
@@ -67,45 +67,31 @@ def __init__(
 
         self.tie_word_embeddings: bool = fd_config.model_config.tie_word_embeddings
 
-        if self.use_ep:
-            self.weight = self.create_parameter(
-                shape=[embedding_dim, num_embeddings],
-                dtype=paddle.get_default_dtype(),
-                is_bias=False,
+        if self.column_cut:
+            need_gather = True
+            self.linear = ColumnParallelLinear(
+                embedding_dim,
+                num_embeddings,
+                mp_group=self.tp_group,
+                weight_attr=None,
+                has_bias=True if self.bias_key is not None else False,
+                gather_output=need_gather,
+                fuse_matmul_bias=False,  # False diff更小
             )
-            if self.bias_key is not None:
-                self.bias = self.create_parameter(
-                    shape=[num_embeddings],
-                    dtype=paddle.get_default_dtype(),
-                    is_bias=True,
-                )
-
+            if self.nranks > 1:
+                set_weight_attrs(self.linear.weight, {"output_dim": True})
         else:
-            if self.column_cut:
-                need_gather = True
-                self.linear = ColumnParallelLinear(
-                    embedding_dim,
-                    num_embeddings,
-                    mp_group=fleet.get_hybrid_communicate_group().get_model_parallel_group(),
-                    weight_attr=None,
-                    has_bias=True if self.bias_key is not None else False,
-                    gather_output=need_gather,
-                    fuse_matmul_bias=False,  # False diff更小
-                )
-                if self.nranks > 1:
-                    set_weight_attrs(self.linear.weight, {"output_dim": True})
-            else:
-                self.linear = RowParallelLinear(
-                    embedding_dim,
-                    num_embeddings,
-                    mp_group=fleet.get_hybrid_communicate_group().get_model_parallel_group(),
-                    weight_attr=None,
-                    has_bias=True if self.bias_key is not None else False,
-                    input_is_parallel=False,
-                    fuse_matmul_bias=False,  # False diff更小
-                )
-                if self.nranks > 1:
-                    set_weight_attrs(self.linear.weight, {"output_dim": False})
+            self.linear = RowParallelLinear(
+                embedding_dim,
+                num_embeddings,
+                mp_group=self.tp_group,
+                weight_attr=None,
+                has_bias=True if self.bias_key is not None else False,
+                input_is_parallel=False,
+                fuse_matmul_bias=False,  # False diff更小
+            )
+            if self.nranks > 1:
+                set_weight_attrs(self.linear.weight, {"output_dim": False})
 
     def load_state_dict(self, state_dict: Dict[str, paddle.Tensor | np.ndarray]):
         """
@@ -115,24 +101,19 @@ def load_state_dict(self, state_dict: Dict[str, paddle.Tensor | np.ndarray]):
             state_dict (dict): A dictionary containing the checkpoint weights and biases.
         """
 
-        if self.use_ep:
-            self.weight.set_value(get_tensor(state_dict.pop(self.weight_key)).astype(paddle.get_default_dtype()))
-            if self.bias_key is not None:
-                self.bias.set_value(get_tensor(state_dict.pop(self.bias_key)).astype(paddle.get_default_dtype()))
+        if self.tie_word_embeddings:
+            self.linear.weight.set_value(
+                get_tensor(state_dict.pop(self.weight_key)).astype(paddle.get_default_dtype()).transpose([1, 0])
+            )
         else:
-            if self.tie_word_embeddings:
-                self.linear.weight.set_value(
-                    get_tensor(state_dict.pop(self.weight_key)).astype(paddle.get_default_dtype()).transpose([1, 0])
-                )
-            else:
-                weight_tensor = get_tensor(state_dict.pop(self.weight_key)).astype(paddle.get_default_dtype())
-                if self.linear.weight.shape != weight_tensor.shape:
-                    weight_tensor = weight_tensor.transpose([1, 0])
-                self.linear.weight.set_value(weight_tensor)
-
-            if self.bias_key is not None:
-                bias = get_tensor(state_dict.pop(self.bias_key)).astype(paddle.get_default_dtype())
-                self.linear.bias.set_value(bias)
+            weight_tensor = get_tensor(state_dict.pop(self.weight_key)).astype(paddle.get_default_dtype())
+            if self.linear.weight.shape != weight_tensor.shape:
+                weight_tensor = weight_tensor.transpose([1, 0])
+            self.linear.weight.set_value(weight_tensor)
+
+        if self.bias_key is not None:
+            bias = get_tensor(state_dict.pop(self.bias_key)).astype(paddle.get_default_dtype())
+            self.linear.bias.set_value(bias)
 
     def forward(self, input: paddle.Tensor) -> paddle.Tensor:
         """
@@ -145,11 +126,5 @@ def forward(self, input: paddle.Tensor) -> paddle.Tensor:
             Tensor: The output tensor after processing through the layer.
         """
         logits = input
-        if self.use_ep:
-            if self.bias_key is None:
-                logits = paddle.matmul(logits, self.weight)
-            else:
-                logits = paddle.incubate.nn.functional.fused_linear(logits, self.weight, self.bias)
-        else:
-            logits = self.linear(logits)
+        logits = self.linear(logits)
         return logits

fastdeploy/model_executor/layers/moe/fused_moe_cutlass_backend.py

@@ -309,7 +309,7 @@ def apply_tp(
         )
 
         if layer.reduce_results and layer.tp_size > 1:
-            tensor_model_parallel_all_reduce(fused_moe_out)
+            tensor_model_parallel_all_reduce(fused_moe_out, self.tp_group)
 
         return fused_moe_out
 

fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py

@@ -465,6 +465,6 @@ def apply_tp(
             1.0,
         )[0]
         if layer.tp_size > 1:
-            tensor_model_parallel_all_reduce(tmp_ffn_out)
+            tensor_model_parallel_all_reduce(tmp_ffn_out, self.tp_group)
 
         return tmp_ffn_out