add some mappings

Author: Difers

paddlenlp/transformers/qwen/modeling.py

@@ -525,15 +525,22 @@ def get_tensor_parallel_split_mappings(num_hidden_layers):
             base_actions = {
                 # Column Linear
                 "lm_head.weight": partial(fn, is_column=True),
-                "qwen.h.0.mlp.w2.weight": partial(fn, is_column=True),
-                "qwen.h.0.mlp.w1.weight": partial(fn, is_column=True),
                 "qwen.h.0.attn.c_attn.weight": partial(fn, is_column=True, is_naive_3fuse=True),
                 "qwen.h.0.attn.c_attn.bias": partial(fn, is_column=True, is_naive_3fuse=True),
                 # Row Linear
                 "qwen.wte.weight": partial(fn, is_column=False),
                 "qwen.h.0.mlp.c_proj.weight": partial(fn, is_column=False),
                 "qwen.h.0.attn.c_proj.weight": partial(fn, is_column=False),
             }
+
+            if config.fuse_attention_ffn:
+                base_actions["layers.0.mlp.gate_up_fused_proj.weight"] = partial(
+                    fn, is_column=True, is_naive_2fuse=True
+                )
+            else:
+                base_actions["qwen.h.0.mlp.w2.weight"] = partial(fn, is_column=True)
+                base_actions["qwen.h.0.mlp.w1.weight"] = partial(fn, is_column=True)
+
             for key, action in base_actions.items():
                 if "h.0." in key:
                     for i in range(num_hidden_layers):
@@ -569,8 +576,8 @@ def _get_name_mappings(cls, config: QWenConfig) -> List[StateDictNameMapping]:
                     f"h.{layer_index}.attn.c_attn.bias",
                 ],
                 [
-                    f"h.{layer_index}.attn.c_proj.weight",
-                    f"h.{layer_index}.attn.c_proj.weight",
+                    f"h.{layer_index}.attn.o_proj.weight",
+                    f"h.{layer_index}.attn.o_proj.weight",
                     "transpose",
                 ],
                 [