Fix bugs and CI and optimize codes

Signed-off-by: zzhx1 <zzh_201018@outlook.com>

Author: zzhx1

tests/ut/models/test_deepseek_v2.py

@@ -27,7 +27,7 @@
     CustomDeepseekV2MLP, CustomDeepseekV2MoE,
     CustomDeepseekV2RowParallelLinear,
     CustomDeepseekV2RowParallelLinearReplaceAllreduce,
-    CustomDeepseekV2SiluAndMul)
+    CustomDeepseekV2SiluAndMul, CustomLogitsProcessor, CustomParallelLMHead)
 
 
 @pytest.fixture
@@ -317,3 +317,19 @@ def test_custom_deepseek_v2_for_causal_lm(mock_distributed, vllm_config):
     ):
         loaded = model.load_weights(weights)
         assert loaded is not None
+
+
+def test_custom_deepseek_v2_lmhead(mock_distributed, vllm_config):
+    model = CustomDeepseekV2ForCausalLM(vllm_config=vllm_config)
+    lmhead = CustomParallelLMHead(model.model.config.vocab_size,
+                                  model.model.config.hidden_size)
+    logits_processor = CustomLogitsProcessor(model.model.config.vocab_size)
+
+    input_ids = torch.randint(0, 10000, (2, 4))
+    positions = torch.arange(4).repeat(2, 1)
+    with patch.object(model.model,
+                      "forward",
+                      return_value=torch.randn(2, 4, 128)):
+        output = model(input_ids, positions)
+        logits = logits_processor(lmhead, output)
+    assert logits.shape == (2, 4, 10000)

vllm_ascend/ascend_config.py

@@ -53,16 +53,15 @@ def __init__(self, vllm_config):
         self.lmhead_tensor_parallel_size = additional_config.get(
             "lmhead_tensor_parallel_size", None)
         if self.lmhead_tensor_parallel_size is not None:
-            logger.info(f"Enable lmhead_tensor_parallel_size={self.lmhead_tensor_parallel_size} in pure DP scenario")
-            assert(
+            logger.info(
+                f"Enable lmhead_tensor_parallel_size={self.lmhead_tensor_parallel_size} in pure DP scenario"
+            )
+            assert (
                 vllm_config.parallel_config.tensor_parallel_size == 1
-            ),"lmhead_tensor_parallel_size is only supported in the pure DP scenario"
-            assert(
-                self.torchair_graph_config.enabled == True
+            ), "lmhead_tensor_parallel_size is only supported in the pure DP scenario"
+            assert (
+                self.torchair_graph_config.enabled
             ), "lmhead_tensor_parallel_size is only supported in graph mode"
-            assert(
-                vllm_config.kv_transfer_config is not None and vllm_config.kv_transfer_config.is_kv_consumer
-            ),"lmhead_tensor_parallel_size is only supported in pd scenario and can only be used in D node."
 
 
 class TorchairGraphConfig:

vllm_ascend/distributed/parallel_state.py

@@ -9,15 +9,18 @@
 _MC2: Optional[GroupCoordinator] = None
 _LMTP: Optional[GroupCoordinator] = None
 
+
 def get_mc2_group() -> GroupCoordinator:
     assert _MC2 is not None, ("mc2 group is not initialized")
     return _MC2
 
+
 def get_lmheadtp_group() -> GroupCoordinator:
     assert _LMTP is not None, (
         "lm head tensor parallel group is not initialized")
     return _LMTP
 
+
 def model_parallel_initialized():
     return (_MC2 is not None)
 
@@ -43,22 +46,23 @@ def init_ascend_model_parallel(parallel_config: ParallelConfig, ):
                                      get_world_group().local_rank,
                                      backend,
                                      group_name="mc2")
-    
+
     lmhead_tensor_parallel_size = parallel_config.lmhead_tensor_parallel_size
     if lmhead_tensor_parallel_size is not None:
         group_ranks = []
         global _LMTP
         num_lmhead_tensor_parallel_groups: int = (world_size //
-                                                lmhead_tensor_parallel_size)
+                                                  lmhead_tensor_parallel_size)
         for i in range(num_lmhead_tensor_parallel_groups):
             ranks = list(
                 range(i * lmhead_tensor_parallel_size,
-                    (i + 1) * lmhead_tensor_parallel_size))
+                      (i + 1) * lmhead_tensor_parallel_size))
             group_ranks.append(ranks)
         _LMTP = init_model_parallel_group(group_ranks,
-                                        get_world_group().local_rank,
-                                        backend,
-                                        group_name="lmheadtp")
+                                          get_world_group().local_rank,
+                                          backend,
+                                          group_name="lmheadtp")
+
 
 def destroy_ascend_model_parallel():
     global _MC2

vllm_ascend/models/deepseek_mtp.py

@@ -25,18 +25,20 @@
 from vllm.attention.backends.abstract import AttentionMetadata
 from vllm.config import CacheConfig, ModelConfig, VllmConfig
 from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.sampler import get_sampler
-from vllm.model_executor.layers.vocab_parallel_embedding import (
-    ParallelLMHead, VocabParallelEmbedding)
+from vllm.model_executor.layers.vocab_parallel_embedding import \
+    VocabParallelEmbedding
 from vllm.model_executor.models.deepseek_mtp import (
     DeepSeekMTP, DeepSeekMultiTokenPredictor, DeepSeekMultiTokenPredictorLayer,
     SharedHead)
 from vllm.model_executor.models.utils import maybe_prefix
 from vllm.model_executor.sampling_metadata import SamplingMetadata
 from vllm.sequence import IntermediateTensors
 
+from vllm_ascend.ops.vocab_parallel_embedding import (CustomLogitsProcessor,
+                                                      CustomParallelLMHead)
+
 from .deepseek_v2 import CustomDeepseekV2DecoderLayer
 
 
@@ -48,10 +50,10 @@ def __init__(self,
                  prefix: str = "") -> None:
         nn.Module.__init__(self)
         self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.head = ParallelLMHead(config.vocab_size,
-                                   config.hidden_size,
-                                   quant_config=quant_config,
-                                   prefix=maybe_prefix(prefix, "head"))
+        self.head = CustomParallelLMHead(config.vocab_size,
+                                         config.hidden_size,
+                                         quant_config=quant_config,
+                                         prefix=maybe_prefix(prefix, "head"))
 
 
 class CustomDeepSeekMultiTokenPredictorLayer(DeepSeekMultiTokenPredictorLayer):
@@ -141,7 +143,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             for idx in range(self.mtp_start_layer_idx,
                              self.mtp_start_layer_idx + self.num_mtp_layers)
         ]
-        self.logits_processor = LogitsProcessor(config.vocab_size)
+        self.logits_processor = CustomLogitsProcessor(config.vocab_size)
 
     def forward(
         self,

vllm_ascend/models/deepseek_v2.py

@@ -49,12 +49,11 @@
                                                ReplicatedLinear,
                                                RowParallelLinear,
                                                UnquantizedLinearMethod)
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.layers.sampler import get_sampler
-from vllm.model_executor.layers.vocab_parallel_embedding import (
-    ParallelLMHead, VocabParallelEmbedding)
+from vllm.model_executor.layers.vocab_parallel_embedding import \
+    VocabParallelEmbedding
 from vllm.model_executor.model_loader.weight_utils import (
     default_weight_loader, maybe_remap_kv_scale_name)
 from vllm.model_executor.models.deepseek_v2 import \
@@ -68,14 +67,15 @@
     PPMissingLayer, is_pp_missing_parameter,
     make_empty_intermediate_tensors_factory, make_layers, maybe_prefix)
 from vllm.sequence import IntermediateTensors
-from vllm.model_executor.sampling_metadata import SamplingMetadata
 
 from vllm_ascend.ascend_config import get_ascend_config
 from vllm_ascend.ops.fused_moe import AscendFusedMoE
+from vllm_ascend.ops.vocab_parallel_embedding import (CustomLogitsProcessor,
+                                                      CustomParallelLMHead)
 from vllm_ascend.quantization.quant_config import AscendLinearMethod
 from vllm_ascend.quantization.w8a8_dynamic import AscendW8A8DynamicLinearMethod
 from vllm_ascend.utils import dispose_tensor, npu_prefetch
-from vllm_ascend.ops.vocab_parallel_embedding import CustomParallelLMHead, CustomLogitsProcessor
+
 
 class CustomDeepseekV2SiluAndMul(SiluAndMul):
 
@@ -930,7 +930,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
                                                 config.hidden_size,
                                                 quant_config=quant_config,
                                                 prefix=maybe_prefix(
-                                                prefix, "lm_head"))
+                                                    prefix, "lm_head"))
         else:
             self.lm_head = PPMissingLayer()
         self.logits_processor = CustomLogitsProcessor(config.vocab_size)

vllm_ascend/ops/vocab_parallel_embedding.py

@@ -19,38 +19,20 @@
 
 import torch
 from torch.nn import Module
-import torch.distributed as dist
-from torch.nn.parameter import Parameter, UninitializedParameter
-
-from vllm.distributed import (
-    divide,
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
-    tensor_model_parallel_all_reduce
-)
-from vllm.model_executor.layers.vocab_parallel_embedding import (
-    VocabParallelEmbedding,
-    DEFAULT_VOCAB_PADDING_SIZE,
-    pad_vocab_size,
-    UnquantizedEmbeddingMethod,
-    ParallelLMHead
-)
-from vllm.model_executor.layers.logits_processor import (
-    LogitsProcessor,
-    _apply_logits_processors,
-    _prune_hidden_states
-)
-from vllm.model_executor.parameter import BasevLLMParameter
-from vllm.model_executor.utils import set_weight_attrs, _enable_lmhead_tp
-from vllm.model_executor.sampling_metadata import SamplingMetadata
+from torch.nn.parameter import Parameter
+from vllm.distributed import (divide, get_tensor_model_parallel_rank,
+                              get_tensor_model_parallel_world_size,
+                              tensor_model_parallel_all_reduce)
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.quantization.base_config import (
-    QuantizationConfig,
-    QuantizeMethodBase,
-    method_has_implemented_embedding
-)
+    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_ascend.distributed.parallel_state import get_lmheadtp_group
-from vllm_ascend.ascend_config import get_ascend_config
+from vllm_ascend.utils import lmhead_tp_enable
 
 
 def get_masked_input_and_mask(
@@ -105,8 +87,7 @@ def vocab_parallel_embedding_forward(self, input_):
 
 
 class CustomParallelLMHead(ParallelLMHead):
-    
-    """Costom Parallelized LM head, added the feature of lmheadTP in pure dp scenario
+    """Custom Parallelized LM head, added the feature of lmheadTP in pure dp scenario
     
     Output logits weight matrices used in the Sampler. The weight and bias
     tensors are padded to make sure they are divisible by the number of
@@ -120,6 +101,7 @@ class CustomParallelLMHead(ParallelLMHead):
         org_num_embeddings: original vocabulary size (without LoRA).
         padding_size: padding size for the vocabulary.
     """
+
     def __init__(self,
                  num_embeddings: int,
                  embedding_dim: int,
@@ -128,16 +110,16 @@ def __init__(self,
                  org_num_embeddings: Optional[int] = None,
                  padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,
                  quant_config: Optional[QuantizationConfig] = None,
-                 prefix: str = ""): 
+                 prefix: str = ""):
         Module.__init__(self)
 
-        if _enable_lmhead_tp():
+        if lmhead_tp_enable():
             tp_rank = get_lmheadtp_group().rank_in_group
             self.tp_size = get_lmheadtp_group().world_size
         else:
             tp_rank = get_tensor_model_parallel_rank()
             self.tp_size = get_tensor_model_parallel_world_size()
-        
+
         self.num_embeddings = num_embeddings
         self.padding_size = padding_size
         self.org_vocab_size = org_num_embeddings or num_embeddings
@@ -197,7 +179,7 @@ def __init__(self,
                                          self.num_embeddings_padded,
                                          params_dtype=params_dtype,
                                          weight_loader=self.weight_loader)
-        
+
         self.quant_config = quant_config
         if bias:
             self.bias = Parameter(
@@ -209,90 +191,32 @@ def __init__(self,
             })
         else:
             self.register_parameter("bias", None)
-        
+
+
 class CustomLogitsProcessor(LogitsProcessor):
     """Custom logits processor extending base LogitsProcessor functionality.
     Added the feature of lmheadTP in pure dp scenario
     """
-    
-    def __init__(self,
-                 vocab_size: int,
-                 org_vocab_size: Optional[int] = None,
-                 scale: float = 1.0,
-                 logits_as_input: bool = False,
-                 soft_cap: Optional[float] = None) -> None:
-        super().__init__(
-            vocab_size=vocab_size,
-            org_vocab_size=org_vocab_size,
-            scale=scale,
-            logits_as_input=logits_as_input,
-            soft_cap=soft_cap
-        )
 
-    def forward(
+    def _get_logits(
         self,
-        lm_head: CustomParallelLMHead,
         hidden_states: torch.Tensor,
-        sampling_metadata: Optional[SamplingMetadata] = None,
-        embedding_bias: Optional[torch.Tensor] = None,
+        lm_head: CustomParallelLMHead,
+        embedding_bias: Optional[torch.Tensor],
     ) -> Optional[torch.Tensor]:
-        if self.logits_as_input:
-            logits = hidden_states
-        else:
-            if sampling_metadata is not None:
-                hidden_states = _prune_hidden_states(hidden_states,
-                                                     sampling_metadata)
-
-            # Get the logits for the next tokens.
-            logits = self._get_logits(hidden_states, lm_head, embedding_bias)
-        if logits is not None:
-            if self.soft_cap is not None:
-                logits = logits / self.soft_cap
-                logits = torch.tanh(logits)
-                logits = logits * self.soft_cap
-
-            if self.scale != 1.0:
-                logits *= self.scale
 
-            # Apply logits processors (if any).
-            if sampling_metadata is not None and \
-                sampling_metadata.seq_groups is not None:
-                logits = _apply_logits_processors(logits, sampling_metadata)
-        
-        return logits
-
-    def _get_logits(
-            self,
-            hidden_states: torch.Tensor,
-            lm_head: CustomParallelLMHead,
-            embedding_bias: Optional[torch.Tensor],
-        ) -> Optional[torch.Tensor]:
-        """
-        Compute logits for next token prediction using parallel processing.
-        
-        Args:
-            hidden_states: Current hidden states from the model with shape [batch_size, hidden_size]
-            lm_head: Parallel embedding layer for vocabulary predictions
-            embedding_bias: Optional bias tensor to add to logits with shape [vocab_size]
-            
-        Returns:
-            Logits tensor for next token prediction with shape [batch_size, vocab_size] or None
-        """
-
-        if _enable_lmhead_tp():
+        if lmhead_tp_enable():
             # Gather hidden states from all devices in tensor parallel group
-            gathered_hidden_states = get_lmheadtp_group().all_gather(hidden_states, dim=0)
+            gathered_hidden_states = get_lmheadtp_group().all_gather(
+                hidden_states, dim=0)
         else:
             gathered_hidden_states = hidden_states
 
-        # Compute logits using quantized matrix multiplication
-        local_logits = lm_head.quant_method.apply(
-            lm_head,
-            gathered_hidden_states,
-            bias=embedding_bias
-        )
+        local_logits = lm_head.quant_method.apply(lm_head,
+                                                  gathered_hidden_states,
+                                                  bias=embedding_bias)
 
-        if _enable_lmhead_tp():
+        if lmhead_tp_enable():
             logits = get_lmheadtp_group().all_to_all(local_logits)
         else:
             # Gather logits for tensor parallel
@@ -301,6 +225,5 @@ def _get_logits(
         # Remove paddings in vocab (if any)
         if logits is not None:
             logits = logits[..., :self.org_vocab_size]
-            
+
         return logits
-