[feat] add custom embedding tensor parallel

Signed-off-by: zzhx1 <[email protected]>

Author: zzh02232027

vllm_ascend/ascend_config.py

@@ -93,6 +93,17 @@ def __init__(self, vllm_config):
                     "oproj_tensor_parallel_size is only supported in pd scenario and can only be used in D node."
                 )
 
+        self.embedding_tensor_parallel_size = additional_config.get(
+            "embedding_tensor_parallel_size", None)
+        if self.embedding_tensor_parallel_size is not None:
+            logger.info(
+                f"Enable embedding_tensor_parallel_size = {self.embedding_tensor_parallel_size} in pure DP scenario"
+            )
+            if vllm_config.parallel_config.tensor_parallel_size != 1:
+                raise AssertionError(
+                    "embedding_tensor_parallel_size is only supported in the pure DP scenario"
+                )
+
 
 class TorchairGraphConfig:
     """

vllm_ascend/distributed/parallel_state.py

@@ -13,6 +13,7 @@
 _MLP_TP: Optional[GroupCoordinator] = None
 _OTP: Optional[GroupCoordinator] = None
 _LMTP: Optional[GroupCoordinator] = None
+_EMBED_TP: Optional[GroupCoordinator] = None
 
 
 def get_mc2_group() -> GroupCoordinator:
@@ -37,6 +38,11 @@ def get_mlp_tp_group() -> GroupCoordinator:
     return _MLP_TP
 
 
+def get_embed_tp_group() -> GroupCoordinator:
+    assert _EMBED_TP is not None, ("emtp group is not initialized")
+    return _EMBED_TP
+
+
 def model_parallel_initialized():
     return (_MC2 is not None)
 
@@ -111,6 +117,23 @@ def init_ascend_model_parallel(parallel_config: ParallelConfig, ):
                                           backend,
                                           group_name="lmheadtp")
 
+    embedding_tensor_parallel_size = get_ascend_config(
+    ).embedding_tensor_parallel_size
+    if embedding_tensor_parallel_size is not None:
+        group_ranks = []
+        global _EMBED_TP
+        num_embedding_tensor_parallel_groups: int = (
+            world_size // embedding_tensor_parallel_size)
+        for i in range(num_embedding_tensor_parallel_groups):
+            ranks = list(
+                range(i * embedding_tensor_parallel_size,
+                      (i + 1) * embedding_tensor_parallel_size))
+            group_ranks.append(ranks)
+        _EMBED_TP = init_model_parallel_group(group_ranks,
+                                              get_world_group().local_rank,
+                                              backend,
+                                              group_name="emtp")
+
 
 def get_mlp_tensor_model_parallel_world_size():
     """Return world size for the tensor model parallel group."""
@@ -142,3 +165,8 @@ def destroy_ascend_model_parallel():
     if _OTP:
         _OTP.destroy()
     _OTP = None
+
+    global _EMBED_TP
+    if _EMBED_TP:
+        _EMBED_TP.destroy()
+    _EMBED_TP = None

vllm_ascend/ops/vocab_parallel_embedding.py

@@ -15,13 +15,16 @@
 # limitations under the License.
 #
 
+from itertools import accumulate
 from typing import Optional, Tuple
 
 import torch
 from torch import nn
 from torch.nn.parameter import Parameter
+from vllm.config import get_current_vllm_config
 from vllm.distributed import divide, tensor_model_parallel_all_reduce
 from vllm.distributed.parallel_state import get_tp_group
+from vllm.forward_context import get_forward_context
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.quantization.base_config import (
     QuantizationConfig, QuantizeMethodBase, method_has_implemented_embedding)
@@ -30,8 +33,10 @@
     VocabParallelEmbedding, pad_vocab_size)
 from vllm.model_executor.utils import set_weight_attrs
 
-from vllm_ascend.distributed.parallel_state import get_lmhead_tp_group
-from vllm_ascend.utils import lmhead_tp_enable
+from vllm_ascend.ascend_config import get_ascend_config
+from vllm_ascend.distributed.parallel_state import (get_embed_tp_group,
+                                                    get_lmhead_tp_group)
+from vllm_ascend.utils import embedding_tp_enable, lmhead_tp_enable
 
 
 class AscendVocabParallelEmbedding(VocabParallelEmbedding):
@@ -51,8 +56,15 @@ def __init__(self,
                  prefix: str = ""):
         nn.Module.__init__(self)
 
+        self.forward_type = None
         if lmhead_tp_enable() and prefix.find("lm_head") != -1:
             self.comm_group = get_lmhead_tp_group()
+        elif embedding_tp_enable() and prefix.find("embed_tokens") != -1:
+            self.comm_group = get_embed_tp_group()
+            self.forward_type = "embed_tp"
+            self.is_decode_only = get_current_vllm_config(
+            ).kv_transfer_config.is_kv_consumer
+            self.forward_type = "embed_tp"
         else:
             self.comm_group = get_tp_group()
 
@@ -146,6 +158,54 @@ def _get_masked_input_and_mask(
         return input_, ~vocab_mask
 
     def forward(self, input_):
+        if self.forward_type == "embed_tp":
+            return self._forward_embed_tp(input_)
+        else:
+            return self._forward_origin(input_)
+
+    def _forward_embed_tp(self, input_):
+        if get_ascend_config(
+        ).torchair_graph_config.enabled is False and not self.is_decode_only:
+            cu_tokens_across_dp_cpu = get_forward_context(
+            ).dp_metadata.cu_tokens_across_dp_cpu
+            global_dp_batch_size = torch.diff(
+                cu_tokens_across_dp_cpu,
+                prepend=cu_tokens_across_dp_cpu.new_zeros(1))
+            embedd_group_batch_size = [
+                global_dp_batch_size[x] for x in self.comm_group.ranks
+            ]
+            # Gather inputs from all embed TP ranks
+            gathered_input = [
+                torch.empty(batch_size, dtype=input_.dtype, device='npu')
+                for batch_size in embedd_group_batch_size
+            ]
+            torch.distributed.all_gather(gathered_input,
+                                         input_,
+                                         group=self.comm_group.device_group)
+            complete_input = torch.cat(gathered_input, dim=0)
+        else:
+            complete_input = self.comm_group.all_gather(input_, dim=0)
+            embedd_group_batch_size = [input_.size(0)
+                                       ] * self.comm_group.world_size
+        # Mask input for vocab sharding
+        masked_input, input_mask = self._get_masked_input_and_mask(
+            complete_input, self.shard_indices.org_vocab_start_index,
+            self.shard_indices.org_vocab_end_index,
+            self.shard_indices.num_org_vocab_padding,
+            self.shard_indices.added_vocab_start_index,
+            self.shard_indices.added_vocab_end_index)
+        complete_output = self.quant_method.embedding(self,
+                                                      masked_input.long())
+        complete_output.masked_fill_(input_mask.unsqueeze(-1), 0)
+        output = self.comm_group.all_reduce(complete_output)
+        # Slice output to return only local batch portion
+        prefix_sum = list(accumulate(embedd_group_batch_size))
+        start_idx = prefix_sum[self.tp_rank - 1] if self.tp_rank > 0 else 0
+        end_idx = prefix_sum[self.tp_rank]
+        output = output[start_idx:end_idx]
+        return output
+
+    def _forward_origin(self, input_):
         if self.tp_size > 1:
             # Build the mask.
             masked_input, input_mask = self._get_masked_input_and_mask(

vllm_ascend/utils.py

@@ -574,6 +574,10 @@ def lmhead_tp_enable() -> bool:
     return get_ascend_config().lmhead_tensor_parallel_size is not None
 
 
+def embedding_tp_enable() -> bool:
+    return get_ascend_config().embedding_tensor_parallel_size is not None
+
+
 def oproj_tp_enable() -> bool:
     return get_ascend_config().oproj_tensor_parallel_size is not None