update TP and model load

Author: DrownFish19

paddlenlp/generation/utils.py

@@ -742,7 +742,7 @@ def generate(
                 # ['是的', '嗯嗯']
         """
         if generation_config is None:
-            if self.generation_config._from_model_config:
+            if self.generation_config is None or self.generation_config._from_model_config:
                 new_generation_config = GenerationConfig.from_model_config(self.config)
                 if new_generation_config != self.generation_config:
                     logger.warning(

paddlenlp/transformers/conversion_utils.py

@@ -1319,7 +1319,7 @@ def _resolve_prefix_keys(state_keys_base, state_keys_real, ignore_error=False):
             for x in state_keys_real:
                 if x.endswith(key):
                     state_keys_map[key] = x
-                    break
+                    # break # remove break for math A.key B.key ...
             if key not in state_keys_map:
                 if not ignore_error:
                     logger.debug(f"tensor parallel conversion: could not find name {key} in loaded state dict!")

paddlenlp/transformers/deepseek_v2/modeling.py

@@ -856,6 +856,9 @@ def __init__(self, config: DeepseekV2Config, layerwise_recompute: bool = False):
 
             self.o_proj = RowParallelLinear(self.num_heads * self.v_head_dim, self.hidden_size, has_bias=config.attention_bias, input_is_parallel=True)
 
+            assert self.num_heads % config.tensor_parallel_degree == 0, f"num_heads: {self.num_heads}, tensor_parallel_degree: {config.tensor_parallel_degree}"
+            self.num_heads = self.num_heads // config.tensor_parallel_degree
+
         else:
             # for without tensor parallel
             if self.q_lora_rank is None:
@@ -1228,12 +1231,15 @@ def get_tensor_parallel_split_mappings(num_layers):
             # Column Linear
             base_actions["layers.0.self_attn.q_proj.weight"] = partial(fn, is_column=True)
             base_actions["layers.0.self_attn.q_proj.bias"] = partial(fn, is_column=True)
+            base_actions["layers.0.self_attn.q_b_proj.weight"] = partial(fn, is_column=True)
+
             # if we have enough num_key_value_heads to split, then split it.
             if config.num_key_value_heads % config.tensor_parallel_degree == 0:
                 base_actions["layers.0.self_attn.k_proj.weight"] = partial(fn, is_column=True)
                 base_actions["layers.0.self_attn.v_proj.weight"] = partial(fn, is_column=True)
                 base_actions["layers.0.self_attn.k_proj.bias"] = partial(fn, is_column=True)
                 base_actions["layers.0.self_attn.v_proj.bias"] = partial(fn, is_column=True)
+                base_actions["layers.0.self_attn.kv_b_proj.weight"] = partial(fn, is_column=True)
 
             base_actions["layers.0.mlp.up_proj.weight"] = partial(fn, is_column=True)
             base_actions["layers.0.mlp.gate_proj.weight"] = partial(fn, is_column=True)
@@ -1625,9 +1631,7 @@ def forward(self, hidden_states, tensor_parallel_output=None):
         if tensor_parallel_output is None:
             tensor_parallel_output = self.config.tensor_parallel_output
 
-        logits = parallel_matmul(
-            hidden_states, self.weight, transpose_y=False, tensor_parallel_output=tensor_parallel_output
-        )
+        logits = parallel_matmul(hidden_states, self.weight, tensor_parallel_output=tensor_parallel_output)
         return logits
 
 
@@ -1639,7 +1643,7 @@ def __init__(self, config: DeepseekV2Config):
         self.config = config
         self.deepseek_v2 = DeepseekV2Model(config)
         self.vocab_size = config.vocab_size
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias_attr=False)
+        self.lm_head = DeepSeekV2LMHead(config)
         self.criterion = DeepSeekV2PretrainingCriterion(config)
 
     def get_input_embeddings(self):

paddlenlp/transformers/deepseek_v3/modeling.py

@@ -24,10 +24,10 @@
 from typing import List, Optional, Tuple, Union
 
 import paddle
-from paddle import nn
 
 from ..deepseek_v2.modeling import (
     DeepseekV2ForSequenceClassification,
+    DeepSeekV2LMHead,
     DeepseekV2Model,
     DeepseekV2PretrainedModel,
     DeepSeekV2PretrainingCriterion,
@@ -63,7 +63,7 @@ def __init__(self, config: DeepseekV2Config):
         super().__init__(config)
         self.deepseek_v3 = DeepseekV3Model(config)
         self.vocab_size = config.vocab_size
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias_attr=False)
+        self.lm_head = DeepSeekV2LMHead(config)
         self.criterion = DeepSeekV2PretrainingCriterion(config)
 
     def get_input_embeddings(self):