[AutoParallel] opt memory usage for tensor-fusion save&load (#70195)

AndSonder · web-flow · commit 301071f9a463 · 2024-12-17T11:06:16.000+08:00
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -30,6 +30,7 @@
 #include "paddle/fluid/framework/executor.h"
 #include "paddle/fluid/framework/ir/pass.h"
 #include "paddle/fluid/framework/program_desc.h"
+#include "paddle/fluid/framework/tensor_util.h"
 #include "paddle/fluid/ir_adaptor/translator/program_translator.h"
 #include "paddle/fluid/ir_adaptor/translator/translate.h"
 #include "paddle/fluid/ir_adaptor/translator/utils.h"
@@ -635,21 +636,32 @@ void BindProgram(py::module *m) {
                   "The mode is not supported."));
             }
           })
-      .def("set_state_dict",
-           [](std::shared_ptr<Program> self,
-              const std::unordered_map<std::string, phi::DenseTensor>
-                  &state_dict,
-              const framework::Scope &scope = framework::Scope()) {
-             for (auto item : state_dict) {
-               auto var = scope.FindVar(item.first);
-               if (var == nullptr) {
-                 PADDLE_THROW(common::errors::NotFound(
-                     "The variable %s is not found.", item.first));
-               } else {
-                 *var->GetMutable<phi::DenseTensor>() = item.second;
-               }
-             }
-           })
+      .def(
+          "set_state_dict",
+          [](std::shared_ptr<Program> self,
+             const std::unordered_map<std::string, phi::DenseTensor>
+                 &state_dict,
+             const framework::Scope &scope = framework::Scope(),
+             bool copy_tensor = false) {
+            for (auto item : state_dict) {
+              auto var = scope.FindVar(item.first);
+              if (var == nullptr) {
+                PADDLE_THROW(common::errors::NotFound(
+                    "The variable %s is not found.", item.first));
+              } else {
+                if (copy_tensor) {
+                  auto *mutable_tensor = var->GetMutable<phi::DenseTensor>();
+                  paddle::framework::TensorCopy(
+                      item.second, item.second.place(), mutable_tensor);
+                } else {
+                  *var->GetMutable<phi::DenseTensor>() = item.second;
+                }
+              }
+            }
+          },
+          py::arg("state_dict"),
+          py::arg("scope"),
+          py::arg("copy_tensor") = false)
       .def(
           "_prune",
           [](Program &self, std::vector<pir::Value> output_vars) {
diff --git a/python/paddle/distributed/auto_parallel/api.py b/python/paddle/distributed/auto_parallel/api.py
@@ -2208,14 +2208,13 @@ def __init__(
             )
             dist.fleet.init(is_collective=True)
 
-        if int(os.environ.get('FLAGS_enable_sharding_stage1_tensor_fusion', 0)):
-            if isinstance(optimizer, _ShardOptimizer) and use_pir_api():
-                shard_fn = optimizer._shard_fn
-                optimizer = optimizer._inner_opt
-                if isinstance(optimizer._shard_fn, ShardingStage1):
-                    optimizer = ShardingOptimizerStage1(
-                        optimizer, shard_fn, self._inner_strategy
-                    )
+        if isinstance(optimizer, _ShardOptimizer) and use_pir_api():
+            shard_fn = optimizer._shard_fn
+            optimizer = optimizer._inner_opt
+            if isinstance(optimizer._shard_fn, ShardingStage1):
+                optimizer = ShardingOptimizerStage1(
+                    optimizer, shard_fn, self._inner_strategy
+                )
 
         self._engine = Engine(
             layer, loss, optimizer, metrics, strategy=self._inner_strategy
@@ -2688,6 +2687,7 @@ def set_state_dict(self, state_dict: dict[str, Tensor]) -> None:
         cur_state_dict = self.state_dict(
             split_fusion=False, load_sharded_model=False
         )
+        copy_tensor = False
 
         # For sharding with tensor-fusion, we need to convert the state_dict
         # to include tensor-fusion parameters before calling set_state_dict,
@@ -2703,6 +2703,11 @@ def set_state_dict(self, state_dict: dict[str, Tensor]) -> None:
                 optimizer.convert_state_dict_with_tensor_fusion_param(
                     state_dict
                 )
+                # When using the tensor-fusion strategy, model parameters are shared with
+                # slice@ parameters. When setting the state_dict, we must copy the tensor
+                # instead of changing the handle directly, as this could cause errors in
+                # the slice@ parameters and increase memory usage.
+                copy_tensor = True
 
         for k, v in state_dict.items():
             assert v.is_dist(), f"key {k} value:{v} is not a dist tensor."
