Refine deepseekv2 modeling for to_static

llm/run_pretrain.py

@@ -179,6 +179,26 @@ class ModelArguments:
         default=None,
         metadata={"help": "num_hidden_layers."},
     )
+    first_k_dense_replace: Optional[int] = field(
+        default=None,
+        metadata={"help": "first_k_dense_replace."},
+    )
+    n_routed_experts: Optional[int] = field(
+        default=None,
+        metadata={"help": "n_routed_experts."},
+    )
+    num_experts_per_tok: Optional[int] = field(
+        default=None,
+        metadata={"help": "num_experts_per_tok."},
+    )
+    hidden_size: Optional[int] = field(
+        default=None,
+        metadata={"help": "hidden_size."},
+    )
+    topk_group: Optional[int] = field(
+        default=None,
+        metadata={"help": "topk_group."},
+    )
 
 
 def create_pretrained_dataset(
@@ -418,6 +438,23 @@ def main():
     config.num_hidden_layers = (
         model_args.num_hidden_layers if model_args.num_hidden_layers is not None else config.num_hidden_layers
     )
+    config.num_hidden_layers = (
+        model_args.num_hidden_layers if model_args.num_hidden_layers is not None else config.num_hidden_layers
+    )
+    config.first_k_dense_replace = (
+        model_args.first_k_dense_replace
+        if model_args.first_k_dense_replace is not None
+        else config.first_k_dense_replace
+    )
+    config.n_routed_experts = (
+        model_args.n_routed_experts if model_args.n_routed_experts is not None else config.n_routed_experts
+    )
+    config.num_experts_per_tok = (
+        model_args.num_experts_per_tok if model_args.num_experts_per_tok is not None else config.num_experts_per_tok
+    )
+    config.hidden_size = model_args.hidden_size if model_args.hidden_size is not None else config.hidden_size
+    config.topk_group = model_args.topk_group if model_args.topk_group is not None else config.topk_group
+
     # Config for model using dropout, such as GPT.
     if hasattr(config, "hidden_dropout_prob"):
         config.hidden_dropout_prob = model_args.hidden_dropout_prob

paddlenlp/transformers/deepseek_v2/modeling.py

@@ -543,7 +543,7 @@
 
         t = paddle.arange(seq_len, dtype=paddle.float32)
 
-        freqs = paddle.outer(t, self.inv_freq)
+        freqs = paddle.outer(t, paddle.cast(self.inv_freq, dtype="float32"))
 
         _mscale = float(
             yarn_get_mscale(self.scaling_factor, self.mscale)

paddlenlp/transformers/deepseek_v2/modeling_auto.py

@@ -64,7 +64,6 @@
     _make_causal_mask,
     apply_rotary_pos_emb,
     get_triangle_upper_mask,
-    is_casual_mask,
     yarn_get_mscale,
 )
 
@@ -705,7 +704,7 @@
             inputs_embeds = self.embed_tokens(input_ids)
 
         # embed positions
-        if attn_mask_startend_row_indices is not None or get_use_casual_mask():
+        if attn_mask_startend_row_indices is not None or get_use_casual_mask() or self.config.use_flash_attention:
             attention_mask = None
         else:
             # [bs, seq_len]
@@ -717,9 +716,6 @@
             attention_mask = self._prepare_decoder_attention_mask(
                 attention_mask, (batch_size, seq_length), past_key_values_length, inputs_embeds.dtype
             )  # [bs, 1, seq_len, seq_len]
-            if self.config.use_flash_attention:
-                attention_mask = None if is_casual_mask(attention_mask) else attention_mask
-
         # embed positions
         hidden_states = inputs_embeds