[CI] fix

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

Author: zzhx1

vllm_ascend/distributed/parallel_state.py

@@ -4,18 +4,19 @@
 from vllm.config import ParallelConfig
 from vllm.distributed.parallel_state import (GroupCoordinator, get_world_group,
                                              init_model_parallel_group)
+from vllm.utils import logger
 
 import vllm_ascend.envs as envs_ascend
 from vllm_ascend.ascend_config import get_ascend_config
 
-from vllm.utils import logger
 # Currently, mc2 op need their own group coordinator.
 _MC2: Optional[GroupCoordinator] = None
 _MLP_TP: Optional[GroupCoordinator] = None
 
 _LMTP: Optional[GroupCoordinator] = None
 _EMTP: Optional[GroupCoordinator] = None
 
+
 def get_mc2_group() -> GroupCoordinator:
     assert _MC2 is not None, ("mc2 group is not initialized")
     return _MC2
@@ -26,10 +27,12 @@ def get_lmhead_tp_group() -> GroupCoordinator:
         "lm head tensor parallel group is not initialized")
     return _LMTP
 
+
 def get_emtp_group() -> GroupCoordinator:
     assert _EMTP is not None, ("emtp group is not initialized")
     return _EMTP
 
+
 def get_mlp_tp_group() -> GroupCoordinator:
     assert _MLP_TP is not None, ("mlp group is not initialized")
     return _MLP_TP
@@ -99,8 +102,8 @@ def init_ascend_model_parallel(parallel_config: ParallelConfig, ):
     if embedding_tensor_parallel_size is not None:
         group_ranks = []
         global _EMTP
-        num_embedding_tensor_parallel_groups: int = (world_size //
-                                                embedding_tensor_parallel_size)
+        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,
@@ -110,7 +113,10 @@ def init_ascend_model_parallel(parallel_config: ParallelConfig, ):
                                           get_world_group().local_rank,
                                           backend,
                                           group_name="emtp")
-        logger.info(f"Successfully established embedding communication parallel group with size {embedding_tensor_parallel_size}")
+        logger.info(
+            f"Successfully established embedding communication parallel group with size {embedding_tensor_parallel_size}"
+        )
+
 
 def get_mlp_tensor_model_parallel_world_size():
     """Return world size for the tensor model parallel group."""

vllm_ascend/ops/vocab_parallel_embedding.py

@@ -21,21 +21,20 @@
 from torch import nn
 from torch.nn.parameter import Parameter
 from vllm.distributed import divide, tensor_model_parallel_all_reduce
-from vllm.distributed.parallel_state import get_tp_group
-import torch.distributed as dist
+from vllm.distributed.parallel_state import get_dp_group, 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)
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, UnquantizedEmbeddingMethod,
     VocabParallelEmbedding, pad_vocab_size)
 from vllm.model_executor.utils import set_weight_attrs
-from vllm.distributed.parallel_state import get_dp_group
-from vllm.forward_context import get_forward_context   
 from vllm.utils import logger
 
-from vllm_ascend.distributed.parallel_state import get_lmhead_tp_group, get_emtp_group
-from vllm_ascend.utils import lmhead_tp_enable, embedding_tp_enable
+from vllm_ascend.distributed.parallel_state import (get_emtp_group,
+                                                    get_lmhead_tp_group)
+from vllm_ascend.utils import embedding_tp_enable, lmhead_tp_enable
 
 
 class AscendVocabParallelEmbedding(VocabParallelEmbedding):
@@ -150,30 +149,41 @@ def _get_masked_input_and_mask(
         input_ = vocab_mask * (input_ - valid_offset)
         return input_, ~vocab_mask
 
-    def _get_local_batch_slice(self, tensor: torch.Tensor, 
-                              batch_sizes: list, 
-                              local_batch_size: int, 
-                              rank: int) -> torch.Tensor:
+    def _get_local_batch_slice(self, tensor: torch.Tensor, batch_sizes: list,
+                               local_batch_size: int,
+                               rank: int) -> torch.Tensor:
         end_idx = batch_sizes[rank]
         start_idx = end_idx - local_batch_size
         return tensor[start_idx:end_idx]
-    
+
     def forward(self, input_):
         if embedding_tp_enable():
-            logger.info(f"rank:{get_dp_group().rank_in_group}  embedding_tp_enable")
+            logger.info(
+                f"rank:{get_dp_group().rank_in_group}  embedding_tp_enable")
             return self._forward_embed_tp(input_)
         else:
             return self._forward_normal(input_)
-        
+
     def _forward_embed_tp(self, input_):
-        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))
-        logger.info(f"debug input_: {input_.shape} \n global_dp_batch_size: {global_dp_batch_size}\n ")
-        lmhead_group_batch_size = [global_dp_batch_size[x] for x in get_lmhead_tp_group().ranks]
+        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))
+        logger.info(
+            f"debug input_: {input_.shape} \n global_dp_batch_size: {global_dp_batch_size}\n "
+        )
+        lmhead_group_batch_size = [
+            global_dp_batch_size[x] for x in get_lmhead_tp_group().ranks
+        ]
         local_batch_size = input_.size(0)
-        gathered_input = [torch.empty(batch_size, dtype=input_.dtype, device='npu') for batch_size in lmhead_group_batch_size]
-        torch.distributed.all_gather(
-            gathered_input, input_, group=get_lmhead_tp_group().device_group)
+        gathered_input = [
+            torch.empty(batch_size, dtype=input_.dtype, device='npu')
+            for batch_size in lmhead_group_batch_size
+        ]
+        torch.distributed.all_gather(gathered_input,
+                                     input_,
+                                     group=get_lmhead_tp_group().device_group)
         complete_input = torch.cat(gathered_input, dim=0)
         masked_input, input_mask = self._get_masked_input_and_mask(
             complete_input, self.shard_indices.org_vocab_start_index,
@@ -182,43 +192,48 @@ def _forward_embed_tp(self, input_):
             self.shard_indices.added_vocab_start_index,
             self.shard_indices.added_vocab_end_index)
         logger.info(f"all_gather_down complete_input: {complete_input.shape}")
-        
+
         output = self.quant_method.embedding(self, masked_input.long())
         output.masked_fill_(input_mask.unsqueeze(-1), 0)
         output = tensor_model_parallel_all_reduce(output)
         #         output = output[lmhead_group_batch_size[get_lmhead_tp_group().rank_in_group]-local_batch_size :lmhead_group_batch_size[get_lmhead_tp_group().rank_in_group]]
         # Extract the local batch portion from the gathered output
         lmhead_tp_group = get_lmhead_tp_group()
-        output = self._get_local_batch_slice(
-            output, 
-            lmhead_group_batch_size, 
-            local_batch_size, 
-            lmhead_tp_group.rank_in_group
-        )
-        logger.info(f"rank:{get_dp_group().rank_in_group}  output: {output.shape}")
+        output = self._get_local_batch_slice(output, lmhead_group_batch_size,
+                                             local_batch_size,
+                                             lmhead_tp_group.rank_in_group)
+        logger.info(
+            f"rank:{get_dp_group().rank_in_group}  output: {output.shape}")
         return output
 
     def _forward_normal(self, input_):
         if self.tp_size > 1:
             # Build the mask.
-            masked_input, input_mask = get_masked_input_and_mask(
+            masked_input, input_mask = self._get_masked_input_and_mask(
                 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)
         else:
             masked_input = input_
-        logger.info(f"rank:{get_dp_group().rank_in_group}  masked_input:{masked_input.shape}")
+        logger.info(
+            f"rank:{get_dp_group().rank_in_group}  masked_input:{masked_input.shape}"
+        )
         # Get the embeddings.
-        output_parallel = self.quant_method.embedding(self, masked_input.long())
-        logger.info(f"rank:{get_dp_group().rank_in_group}  output_parallel:{output_parallel.shape}")
+        output_parallel = self.quant_method.embedding(self,
+                                                      masked_input.long())
+        logger.info(
+            f"rank:{get_dp_group().rank_in_group}  output_parallel:{output_parallel.shape}"
+        )
         # Mask the output embedding.
         if self.tp_size > 1:
             output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0)
         # Reduce across all the model parallel GPUs.
         output = tensor_model_parallel_all_reduce(output_parallel)
-        logger.info(f"rank:{get_dp_group().rank_in_group}  forward_normal output:{output.shape}")
+        logger.info(
+            f"rank:{get_dp_group().rank_in_group}  forward_normal output:{output.shape}"
+        )
         return output
 
 

vllm_ascend/utils.py

@@ -557,5 +557,6 @@ def get_ascend_soc_version():
 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
+    return get_ascend_config().embedding_tensor_parallel_size is not None

vllm_ascend/worker/model_runner_v1.py

@@ -90,8 +90,7 @@
 from vllm_ascend.torchair.torchair_mla import AscendMLATorchairMetadata
 from vllm_ascend.utils import (ACL_FORMAT_FRACTAL_ND, ACL_FORMAT_FRACTAL_NZ,
                                ProfileExecuteDuration, is_310p,
-                               lmhead_tp_enable, vllm_version_is,
-                               embedding_tp_enable)
+                               lmhead_tp_enable, vllm_version_is)
 from vllm_ascend.worker.eagle_proposer_v1 import EagleProposer
 from vllm_ascend.worker.mtp_proposer_v1 import MtpProposer
 from vllm_ascend.worker.npu_input_batch import CachedRequestState, InputBatch