[RL] fix qwen load when fuse is enabled and modify gate presion to fp32 (#10842)

Author: SylarTiaNII

paddlenlp/transformers/qwen2_moe/modeling.py

@@ -765,7 +765,7 @@ def __init__(self, config, num_experts, expert_hidden_size, **kwargs):
         # [hidden_size, n_expert]
         self.weight = paddle.create_parameter(
             shape=[expert_hidden_size, num_experts],
-            dtype=paddle.get_default_dtype(),
+            dtype="float32",
             is_bias=False,
             default_initializer=nn.initializer.Constant(1.0),
         )
@@ -1031,14 +1031,18 @@ def get_tensor_parallel_split_mappings(num_layers, num_experts):
                 base_actions.pop("embed_tokens.weight")
 
             # 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)
-            # 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)
+            if config.fuse_attention_qkv:
+                base_actions["layers.0.self_attn.qkv_proj.weight"] = partial(fn, is_column=True)
+                base_actions["layers.0.self_attn.qkv_proj.bias"] = partial(fn, is_column=True)
+            else:
+                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)
+                # 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)
 
             for key, action in base_actions.items():
                 if "layers.0." in key:
@@ -1047,11 +1051,20 @@ def get_tensor_parallel_split_mappings(num_layers, num_experts):
                 final_actions[key] = action
 
             # Add tp split for expert params.
-            base_actions = {
-                "layers.0.mlp.experts.0.gate_proj.weight": partial(fn, is_column=True),
-                "layers.0.mlp.experts.0.down_proj.weight": partial(fn, is_column=False),
-                "layers.0.mlp.experts.0.up_proj.weight": partial(fn, is_column=True),
-            }
+            if config.fuse_attention_ffn:
+                base_actions = {
+                    "layers.0.mlp.experts.0.gate_up_fused_proj.weight": partial(
+                        fn, is_column=True, is_naive_2fuse=True
+                    ),
+                    "layers.0.mlp.experts.0.down_proj.weight": partial(fn, is_column=False),
+                }
+            else:
+                # Add tp split for expert params.
+                base_actions = {
+                    "layers.0.mlp.experts.0.gate_proj.weight": partial(fn, is_column=True),
+                    "layers.0.mlp.experts.0.up_proj.weight": partial(fn, is_column=True),
+                    "layers.0.mlp.experts.0.down_proj.weight": partial(fn, is_column=False),
+                }
             for key, action in base_actions.items():
                 for i in range(num_layers):
                     newkey = key.replace("layers.0.", f"layers.{i}.")
@@ -1060,11 +1073,19 @@ def get_tensor_parallel_split_mappings(num_layers, num_experts):
                         final_actions[newkey2] = action
 
             # Add tp split for shared expert params.
-            base_actions = {
-                "layers.0.mlp.shared_expert.gate_proj.weight": partial(fn, is_column=True),
-                "layers.0.mlp.shared_expert.up_proj.weight": partial(fn, is_column=True),
-                "layers.0.mlp.shared_expert.down_proj.weight": partial(fn, is_column=False),
-            }
+            if config.fuse_attention_ffn:
+                base_actions = {
+                    "layers.0.mlp.shared_expert.gate_up_fused_proj.weight": partial(
+                        fn, is_column=True, is_naive_2fuse=True
+                    ),
+                    "layers.0.mlp.shared_expert.down_proj.weight": partial(fn, is_column=False),
+                }
+            else:
+                base_actions = {
+                    "layers.0.mlp.shared_expert.gate_proj.weight": partial(fn, is_column=True),
+                    "layers.0.mlp.shared_expert.up_proj.weight": partial(fn, is_column=True),
+                    "layers.0.mlp.shared_expert.down_proj.weight": partial(fn, is_column=False),
+                }
             for key, action in base_actions.items():
                 if "layers.0." in key:
                     for i in range(num_layers):
@@ -1101,24 +1122,24 @@ def _get_fuse_or_split_param_mappings(cls, config: Qwen2MoeConfig, is_fuse=False
         ]
 
         fuse_gate_up_keys = (
-            "layers.0.mlp.gate_proj.weight",
-            "layers.0.mlp.up_proj.weight",
-            "layers.0.mlp.gate_up_fused_proj.weight",
+            "layers.0.mlp.experts.0.gate_proj.weight",
+            "layers.0.mlp.experts.0.up_proj.weight",
+            "layers.0.mlp.experts.0.gate_up_fused_proj.weight",
         )
         num_heads = config.num_attention_heads
         num_key_value_heads = getattr(config, "num_key_value_heads", num_heads)
         fuse_attention_qkv = getattr(config, "fuse_attention_qkv", False)
         fuse_attention_ffn = getattr(config, "fuse_attention_ffn", False)
+        num_experts = getattr(config, "num_experts", 128)
 
         final_actions = {}
         if is_fuse:
             if fuse_attention_qkv:
                 for i in range(config.num_hidden_layers):
-                    for fuse_keys in fuse_qkv_keys:
-                        keys = tuple([key.replace("layers.0.", f"layers.{i}.") for key in fuse_keys])
-                        final_actions[keys] = partial(
-                            fn, is_qkv=True, num_heads=num_heads, num_key_value_heads=num_key_value_heads
-                        )
+                    keys = [key.replace("layers.0.", f"layers.{i}.") for key in fuse_gate_up_keys]
+                    for j in range(num_experts):
+                        experts_keys = tuple([key.replace("experts.0.", f"experts.{j}.") for key in keys])
+                        final_actions[experts_keys] = fn
             if fuse_attention_ffn:
                 for i in range(config.num_hidden_layers):
                     keys = tuple([key.replace("layers.0.", f"layers.{i}.") for key in fuse_gate_up_keys])
@@ -1137,8 +1158,10 @@ def _get_fuse_or_split_param_mappings(cls, config: Qwen2MoeConfig, is_fuse=False
                         )
             if not fuse_attention_ffn:
                 for i in range(config.num_hidden_layers):
-                    keys = tuple([key.replace("layers.0.", f"layers.{i}.") for key in fuse_gate_up_keys])
-                    final_actions[keys] = partial(fn, split_nums=2)
+                    keys = [key.replace("layers.0.", f"layers.{i}.") for key in fuse_gate_up_keys]
+                    for j in range(num_experts):
+                        experts_keys = tuple([key.replace("experts.0.", f"experts.{j}.") for key in keys])
+                        final_actions[experts_keys] = partial(fn, split_nums=2)
         return final_actions
 
     def _init_weights(self, layer):

paddlenlp/transformers/qwen3_moe/modeling.py

@@ -414,11 +414,18 @@ def get_tensor_parallel_split_mappings(num_layers, num_experts):
                 base_actions.pop("embed_tokens.weight")
 
             # Column Linear
-            base_actions["layers.0.self_attn.q_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)
+            if config.fuse_attention_qkv:
+                base_actions["layers.0.self_attn.qkv_proj.weight"] = partial(fn, is_column=True)
+                base_actions["layers.0.self_attn.qkv_proj.bias"] = partial(fn, is_column=True)
+            else:
+                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)
+                # 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)
 
             for key, action in base_actions.items():
                 if "layers.0." in key:
@@ -427,11 +434,20 @@ def get_tensor_parallel_split_mappings(num_layers, num_experts):
                 final_actions[key] = action
 
             # Add tp split for expert params.
-            base_actions = {
-                "layers.0.mlp.experts.0.gate_proj.weight": partial(fn, is_column=True),
-                "layers.0.mlp.experts.0.down_proj.weight": partial(fn, is_column=False),
-                "layers.0.mlp.experts.0.up_proj.weight": partial(fn, is_column=True),
-            }
+            if config.fuse_attention_ffn:
+                base_actions = {
+                    "layers.0.mlp.experts.0.gate_up_fused_proj.weight": partial(
+                        fn, is_column=True, is_naive_2fuse=True
+                    ),
+                    "layers.0.mlp.experts.0.down_proj.weight": partial(fn, is_column=False),
+                }
+            else:
+                # Add tp split for expert params.
+                base_actions = {
+                    "layers.0.mlp.experts.0.gate_proj.weight": partial(fn, is_column=True),
+                    "layers.0.mlp.experts.0.up_proj.weight": partial(fn, is_column=True),
+                    "layers.0.mlp.experts.0.down_proj.weight": partial(fn, is_column=False),
+                }
             for key, action in base_actions.items():
                 for i in range(num_layers):
                     newkey = key.replace("layers.0.", f"layers.{i}.")
@@ -471,14 +487,15 @@ def _get_fuse_or_split_param_mappings(cls, config: Qwen3MoeConfig, is_fuse=False
         ]
 
         fuse_gate_up_keys = (
-            "layers.0.mlp.gate_proj.weight",
-            "layers.0.mlp.up_proj.weight",
-            "layers.0.mlp.gate_up_fused_proj.weight",
+            "layers.0.mlp.experts.0.gate_proj.weight",
+            "layers.0.mlp.experts.0.up_proj.weight",
+            "layers.0.mlp.experts.0.gate_up_fused_proj.weight",
         )
         num_heads = config.num_attention_heads
         num_key_value_heads = getattr(config, "num_key_value_heads", num_heads)
         fuse_attention_qkv = getattr(config, "fuse_attention_qkv", False)
         fuse_attention_ffn = getattr(config, "fuse_attention_ffn", False)
+        num_experts = getattr(config, "num_experts", 128)
 
         final_actions = {}
         if is_fuse:
@@ -491,8 +508,10 @@ def _get_fuse_or_split_param_mappings(cls, config: Qwen3MoeConfig, is_fuse=False
                         )
             if fuse_attention_ffn:
                 for i in range(config.num_hidden_layers):
-                    keys = tuple([key.replace("layers.0.", f"layers.{i}.") for key in fuse_gate_up_keys])
-                    final_actions[keys] = fn
+                    keys = [key.replace("layers.0.", f"layers.{i}.") for key in fuse_gate_up_keys]
+                    for j in range(num_experts):
+                        experts_keys = tuple([key.replace("experts.0.", f"experts.{j}.") for key in keys])
+                        final_actions[experts_keys] = fn
         else:
             if not fuse_attention_qkv:
                 for i in range(config.num_hidden_layers):
@@ -507,8 +526,10 @@ def _get_fuse_or_split_param_mappings(cls, config: Qwen3MoeConfig, is_fuse=False
                         )
             if not fuse_attention_ffn:
                 for i in range(config.num_hidden_layers):
-                    keys = tuple([key.replace("layers.0.", f"layers.{i}.") for key in fuse_gate_up_keys])
-                    final_actions[keys] = partial(fn, split_nums=2)
+                    keys = [key.replace("layers.0.", f"layers.{i}.") for key in fuse_gate_up_keys]
+                    for j in range(num_experts):
+                        experts_keys = tuple([key.replace("experts.0.", f"experts.{j}.") for key in keys])
+                        final_actions[experts_keys] = partial(fn, split_nums=2)
         return final_actions
 
     def _init_weights(self, layer):