PaddlePaddle
diff --git a/‎paddle/fluid/pybind/pir.cc
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions b/‎paddle/fluid/pybind/pir.cc
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions
diff --git a/‎python/paddle/distributed/auto_parallel/api.py
Copy file name to clipboard
+127-2Lines changed: 127 additions & 2 deletions b/‎python/paddle/distributed/auto_parallel/api.py
Copy file name to clipboard
+127-2Lines changed: 127 additions & 2 deletions
diff --git a/‎python/paddle/distributed/auto_parallel/static/engine.py
Copy file name to clipboard
+16Lines changed: 16 additions & 0 deletions b/‎python/paddle/distributed/auto_parallel/static/engine.py
Copy file name to clipboard
+16Lines changed: 16 additions & 0 deletions
@@ -853,6 +853,12 @@ void BindOperation(py::module *m) {
              }
              return attrs_dict;
            })
+      .def("copy_attrs_from",
+           [](Operation &self, Operation &other) {
+             for (auto &pair : other.attributes()) {
+               self.set_attribute(pair.first, pair.second);
+             }
+           })
       .def("set_execution_stream",
            [](Operation &self, const std::string &exe_stream) {
              self.set_attribute(
 
@@ -14,6 +14,7 @@
 from __future__ import annotations
 
 import copy
+import warnings
 from types import MethodType
 from typing import TYPE_CHECKING, Any, Literal, TypedDict
 
@@ -47,7 +48,9 @@
 from paddle.distributed.auto_parallel.static.dist_op import DistributedOperator
 from paddle.distributed.auto_parallel.static.utils import (
     convert_to_dims_mapping,
+    fuse_param_func,
     get_dist_attr,
+    split_param_func,
     to_list,
 )
 from paddle.framework import core
@@ -130,6 +133,13 @@ def _to_lodtensor(tensor: paddle.Tensor):
     return lodtensor
 
 
+def _get_suffix(s, prefix):
+    if s.startswith(prefix):
+        return s[len(prefix) :]
+    else:
+        return None
+
+
 class DistAttr(core.TensorDistAttr):
     """
     DistAttr specifies how tensors are distributed or sliced on ProcessMesh.
@@ -2456,7 +2466,9 @@ def _fetch_value(self, value, name=None):
         self._engine._pir_user_defined_fetch_names.append(name)
 
     def state_dict(
-        self, mode: Literal['opt', 'param', 'all'] = "all"
+        self,
+        mode: Literal['opt', 'param', 'all'] = "all",
+        split_fusion: bool = True,
     ) -> dict[str, Tensor]:
         """
         Get the state dict of model and optimizer.
@@ -2481,6 +2493,60 @@ def state_dict(
 
         dist_state_dict = self._build_distributed_state_dict(local_state_dict)
 
+        # The parameters fused in the ffn and qkv fusion pass will be split back into their original, unfused state.
+        if self._engine.fused_ffn_qkv is not None and split_fusion:
+            with paddle.base.dygraph.guard():
+                # Traverse each fusion structure, the key could be ffn or qkv.
+                for key, pat_list in self._engine.fused_ffn_qkv.items():
+                    # Traverse each fusion pattern dict, such as: fused_p1_p2:[p1, p2].
+                    for fusion_map in pat_list:
+                        ((fused_param, ori_params_meta),) = fusion_map.items()
+                        origin_params = list(dist_state_dict.keys())
+                        for param in origin_params:
+                            suffix = _get_suffix(param, fused_param)
+                            if suffix is not None:
+                                value = dist_state_dict[param]
+                                assert (
+                                    value.is_dist()
+                                ), f"key {param} value:{value} is not a dist tensor."
+                                mesh = value.process_mesh
+                                placements = value.placements
+                                if "_pow_acc" in suffix:
+                                    out = (value._local_value(),) * len(
+                                        ori_params_meta
+                                    )
+                                else:
+                                    if len(ori_params_meta) == 3:
+                                        is_qkv = True
+                                        num_heads = ori_params_meta[
+                                            0
+                                        ].local_num_head
+                                        num_key_value_heads = ori_params_meta[
+                                            1
+                                        ].local_num_head
+                                    else:
+                                        is_qkv = False
+                                        num_heads = None
+                                        num_key_value_heads = None
+                                    out = split_param_func(
+                                        value._local_value(),
+                                        split_nums=len(ori_params_meta),
+                                        is_qkv=is_qkv,
+                                        num_heads=num_heads,
+                                        num_key_value_heads=num_key_value_heads,
+                                    )
+                                for i in range(len(ori_params_meta)):
+                                    dist_tensor = dtensor_from_local(
+                                        out[i], mesh, placements
+                                    )
+                                    paddle.assign(
+                                        out[i], dist_tensor._local_value()
+                                    )
+                                    dist_state_dict[
+                                        ori_params_meta[i].name + suffix
+                                    ] = dist_tensor
+                                dist_state_dict.pop(param)
+
         mapping_names = [
             (
                 self._parameter_to_structured_name[k]
@@ -2548,7 +2614,7 @@ def build_distributed_tensor(local_tensor, dist_attr):
     def set_state_dict(self, state_dict: dict[str, Tensor]) -> None:
         local_state_dict = {}
         dist_main_program = self.dist_main_program(mode=self._engine._mode)
-        cur_state_dict = self.state_dict()
+        cur_state_dict = self.state_dict(split_fusion=False)
         for k, v in state_dict.items():
             assert v.is_dist(), f"key {k} value:{v} is not a dist tensor."
             if k in cur_state_dict:
@@ -2564,6 +2630,65 @@ def set_state_dict(self, state_dict: dict[str, Tensor]) -> None:
                 else k
             )
             local_state_dict[param_name] = _to_lodtensor(v._local_value())
+
+        # The structure of ffn and qkv in the network has been fused, and the unfused parameters in the original state_dict are fused.
+        if self._engine.fused_ffn_qkv is not None:
+            with paddle.base.dygraph.guard():
+                # Traverse each fusion structure, the key could be ffn or qkv.
+                for key, pat_list in self._engine.fused_ffn_qkv.items():
+                    # Traverse each fusion pattern dict, such as: fused_p1_p2:[p1, p2].
+                    for fusion_map in pat_list:
+                        ((fused_param, ori_params_meta),) = fusion_map.items()
+                        # Obtain all the parameters to be fused, differentiated by suffixes, such as: beta1_pow_acc_0, _fp32_master_0_moment1_0.
+                        suffix_names = []
+                        for k, v in local_state_dict.items():
+                            suffix = _get_suffix(ori_params_meta[0].name, k)
+                            if suffix is not None:
+                                suffix_names.append(suffix)
+                        if len(suffix_names) == 0:
+                            continue
+                        # Traverse through each parameter for fusion, insert the fused parameters, and delete the pre-fusion parameters.
+                        for suffix in suffix_names:
+                            concat_tensors = []
+                            for ori_p in ori_params_meta:
+                                if ori_p.name + suffix not in local_state_dict:
+                                    warnings.warn(
+                                        f"{ori_p.name + suffix} is not in state_dict."
+                                    )
+                                    break
+                                else:
+                                    concat_tensors.append(
+                                        local_state_dict[ori_p.name + suffix]
+                                    )
+                            if len(concat_tensors) == len(ori_params_meta):
+                                if "_pow_acc" in suffix:
+                                    fused_w = concat_tensors[0]
+                                else:
+                                    if len(ori_params_meta) == 3:
+                                        is_qkv = True
+                                        num_heads = ori_params_meta[
+                                            0
+                                        ].local_num_head
+                                        num_key_value_heads = ori_params_meta[
+                                            1
+                                        ].local_num_head
+                                    else:
+                                        is_qkv = False
+                                        num_heads = None
+                                        num_key_value_heads = None
+                                    fused_w = fuse_param_func(
+                                        concat_tensors,
+                                        is_qkv=is_qkv,
+                                        num_heads=num_heads,
+                                        num_key_value_heads=num_key_value_heads,
+                                    )
+
+                                local_state_dict[fused_param + suffix] = (
+                                    _to_lodtensor(fused_w)
+                                )
+                                for ori_p in ori_params_meta:
+                                    local_state_dict.pop(ori_p + suffix)
+
         dist_main_program.set_state_dict(
             local_state_dict, paddle.static.global_scope()
         )
 
@@ -71,6 +71,7 @@
     apply_partition_pass,
     check_chunk_id,
     complete_chunk_id,
+    fuse_attention_ffn_qkv_pass,
     pipeline_pass,
     remove_unuseful_comm_op_pass,
 )
@@ -331,6 +332,8 @@ def __init__(
 
         self.enable_job_schedule_profiler = False
 
+        self.fused_ffn_qkv = None
+
     # get dist input spec from shard dataloader
     def _prepare_data_spec_from_dataloader(self, dataloader):
         inputs_spec = []
@@ -702,6 +705,19 @@ def _parallel_pir(self, mode):
         mix_fw_program = self._fwd_main_progs[mode]
         startup_program = self._startup_progs[mode]
 
+        # TODO(zhangbo) Open fused_ffn/fused_attention_qkv pass
+        if os.getenv("FLAGS_enable_fused_ffn_qkv_pass") in [
+            'True',
+            'true',
+            '1',
+        ]:
+            self.fused_ffn_qkv = fuse_attention_ffn_qkv_pass(
+                startup_program,
+                mix_fw_program,
+                self.concrete_program,
+                mode="all",
+            )
+
         forward_op_start_idx = 0
         backward_op_start_idx = -1
         opt_op_start_idx = -1