fix model run error when use auto parallel and recompute(use_reentrant=false) (#65188)

jeff41404 · web-flow · commit 86206c8f79f8 · 2024-06-20T18:17:22.000+08:00
* fix model run error when auto parallel and recompute and use_reentrant=false

* solve the defect of TensorWrapper not considering DistTensor

* add unittest

* fix recompute have not support cpu when use_reentrant is false
diff --git a/paddle/fluid/eager/tensor_wrapper.h b/paddle/fluid/eager/tensor_wrapper.h
@@ -152,16 +152,36 @@ class TensorWrapper {
 #ifndef PADDLE_NO_PYTHON
     if (packed_value_ && unpack_hook_) {
       auto tensor_unpacked = (*unpack_hook_)(packed_value_);
-      auto src_dense_tensor =
-          static_cast<phi::DenseTensor*>(tensor_unpacked.impl().get());
+      phi::DenseTensor* src_dense_tensor = nullptr;
+      if (tensor_unpacked.is_dense_tensor()) {
+        VLOG(6) << "tensor_unpacked is DenseTensor";
+        src_dense_tensor =
+            static_cast<phi::DenseTensor*>(tensor_unpacked.impl().get());
+      } else if (tensor_unpacked.is_dist_tensor()) {
+        VLOG(6) << "tensor_unpacked is DistTensor";
+        src_dense_tensor = static_cast<phi::distributed::DistTensor*>(
+                               tensor_unpacked.impl().get())
+                               ->unsafe_mutable_value();
+      } else {
+        PADDLE_THROW(
+            paddle::platform::errors::Fatal("Unrecognized tensor_unpacked type "
+                                            "for egr::TensorWrapper::recover"));
+      }
+
       if (intermidiate_tensor_.is_dense_tensor()) {
+        VLOG(6) << "intermidiate_tensor_ is DenseTensor";
         static_cast<phi::DenseTensor*>(intermidiate_tensor_.impl().get())
             ->ResetHolder(src_dense_tensor->MoveMemoryHolder());
       } else if (intermidiate_tensor_.is_dist_tensor()) {
+        VLOG(6) << "intermidiate_tensor_ is DistTensor";
         static_cast<phi::distributed::DistTensor*>(
             intermidiate_tensor_.impl().get())
             ->unsafe_mutable_value()
             ->ResetHolder(src_dense_tensor->MoveMemoryHolder());
+      } else {
+        PADDLE_THROW(paddle::platform::errors::Fatal(
+            "Unrecognized intermidiate_tensor_ type for "
+            "egr::TensorWrapper::recover"));
       }
     } else {
 #endif
diff --git a/paddle/fluid/pybind/eager.cc b/paddle/fluid/pybind/eager.cc
@@ -47,10 +47,12 @@ limitations under the License. */
 #include "paddle/phi/core/string_tensor.h"
 
 using phi::distributed::DistTensor;
+using phi::distributed::DistTensorMeta;
 using phi::distributed::Placement;
 using phi::distributed::Placements;
 using phi::distributed::ProcessMesh;
 using phi::distributed::TensorDistAttr;
+using phi::distributed::auto_parallel::str_join;
 
 namespace paddle {
 namespace pybind {
@@ -81,35 +83,51 @@ void EmptyTensorInitializer(TensorObject* self,
                                 paddle::framework::proto::VarType::FP32,
                             const std::vector<int>& dims = {0},
                             framework::proto::VarType::Type var_type =
-                                paddle::framework::proto::VarType::LOD_TENSOR) {
+                                paddle::framework::proto::VarType::LOD_TENSOR,
+                            ProcessMesh* process_mesh = nullptr,
+                            Placements* placements = nullptr) {
   auto ddims = common::make_ddim(dims);
   self->tensor.set_name(name);
   auto autograd_meta = egr::EagerUtils::autograd_meta(&(self->tensor));
   autograd_meta->SetPersistable(persistable);
   if (stop_gradient != -1) {
     autograd_meta->SetStopGradient(static_cast<bool>(stop_gradient));
   }
-  if (var_type == paddle::framework::proto::VarType::LOD_TENSOR) {
-    // TODO(jiabin): Maybe support LOD later
-    std::shared_ptr<phi::DenseTensor> dense_tensor = nullptr;
-    if (dims.size() == 1 && dims[0] == 0) {
-      std::shared_ptr<phi::Allocation> allocation_ptr = nullptr;
-      dense_tensor = std::make_shared<phi::DenseTensor>(
-          allocation_ptr,
-          phi::DenseTensorMeta(paddle::framework::TransToPhiDataType(dtype),
-                               ddims));
-    } else {
-      // TODO(dev): we need enhance check for ddims.
-      dense_tensor = std::make_shared<phi::DenseTensor>(
-          std::make_shared<phi::Allocation>(),
-          phi::DenseTensorMeta(paddle::framework::TransToPhiDataType(dtype),
-                               ddims));
+  if (process_mesh != nullptr) {
+#ifdef PADDLE_WITH_DISTRIBUTE
+    VLOG(6) << "in EmptyTensorInitializer, create DistTensor";
+    self->tensor.set_impl(std::make_shared<DistTensor>());
+#else
+    PADDLE_THROW(platform::errors::Unavailable(
+        "The tensor-based initialization of (Dist)Tensor is not supported "
+        "in the current PaddlePaddle, please recompile and install "
+        "PaddlePaddle "
+        "with the option of `WITH_DISTRIBUTE=ON`."));
+#endif
+  } else {
+    VLOG(6) << "in EmptyTensorInitializer, create DenseTensor";
+    if (var_type == paddle::framework::proto::VarType::LOD_TENSOR) {
+      // TODO(jiabin): Maybe support LOD later
+      std::shared_ptr<phi::DenseTensor> dense_tensor = nullptr;
+      if (dims.size() == 1 && dims[0] == 0) {
+        std::shared_ptr<phi::Allocation> allocation_ptr = nullptr;
+        dense_tensor = std::make_shared<phi::DenseTensor>(
+            allocation_ptr,
+            phi::DenseTensorMeta(paddle::framework::TransToPhiDataType(dtype),
+                                 ddims));
+      } else {
+        // TODO(dev): we need enhance check for ddims.
+        dense_tensor = std::make_shared<phi::DenseTensor>(
+            std::make_shared<phi::Allocation>(),
+            phi::DenseTensorMeta(paddle::framework::TransToPhiDataType(dtype),
+                                 ddims));
+      }
+      self->tensor.set_impl(dense_tensor);
+    } else if (var_type == paddle::framework::proto::VarType::SELECTED_ROWS) {
+      std::shared_ptr<phi::SelectedRows> tensor =
+          std::make_shared<phi::SelectedRows>();
+      self->tensor.set_impl(tensor);
     }
-    self->tensor.set_impl(dense_tensor);
-  } else if (var_type == paddle::framework::proto::VarType::SELECTED_ROWS) {
-    std::shared_ptr<phi::SelectedRows> tensor =
-        std::make_shared<phi::SelectedRows>();
-    self->tensor.set_impl(tensor);
   }
 
   if (!autograd_meta->GetMutableGradNode()) {
@@ -768,12 +786,16 @@ Tensor is the basic data structure in PaddlePaddle. There are some ways to creat
  * 1.
  * def __init__ ()
  * 2.
+ * (should have at least five parameter, five parameters create DenseTensor,
+ * seven parameters create DistTensor)
  * def __init__ (
  * ** dtype: paddle::framework::proto::VarType::Type,
  * ** dims: vector<int>,
  * ** name: std::string,
  * ** type: paddle::framework::proto::VarType::LodTensor,
- * ** persistable: bool)
+ * ** persistable: bool,
+ * ** process_mesh: phi::distributed::ProcessMesh,
+ * ** placements: std::vector<Placement>)
  * 3. (multi-place)
  * (should have at least one parameter, one parameter equals to case 4, zero
  * parameter equals to case 1)
@@ -797,7 +819,7 @@ Tensor is the basic data structure in PaddlePaddle. There are some ways to creat
  * ** global_tensor: Tensor,
  * ** place: paddle::platform::Place,
  * ** name: std::string,
- * ** process_mesh: phi::distributed::ProcessMesh)
+ * ** process_mesh: phi::distributed::ProcessMesh,
  * ** placements: std::vector<Placement>)
  * 7. (multi-place)
  * (should have at least one parameter, one parameter equals to case 5, zero
@@ -806,7 +828,7 @@ Tensor is the basic data structure in PaddlePaddle. There are some ways to creat
  * ** local_tensor: Tensor,
  * ** global_dims: vector<int>,
  * ** name: std::string,
- * ** process_mesh: phi::distributed::ProcessMesh)
+ * ** process_mesh: phi::distributed::ProcessMesh,
  * ** placements: std::vector<Placement>)
  * 8. (multi-place) (should have at least one parameter, one parameter similar
  * to case 5, zero parameter equals to case 1.)
@@ -995,14 +1017,28 @@ int TensorInit(PyObject* self, PyObject* args, PyObject* kwargs) {
             CastPyArg2ProtoType(kw_type, 0);
         bool persistable = CastPyArg2AttrBoolean(kw_persistable, 0);
 
+        ProcessMesh* process_mesh_ptr = nullptr;
+        if (kw_process_mesh != nullptr) {
+          ProcessMesh process_mesh = CastPyArg2ProcessMesh(kw_process_mesh, 0);
+          process_mesh_ptr = &process_mesh;
+        }
+
+        Placements* placements_ptr = nullptr;
+        if (kw_placements != nullptr) {
+          Placements placements = CastPyArg2VectorOfPlacement(kw_placements, 0);
+          placements_ptr = &placements;
+        }
+
         EmptyTensorInitializer(py_tensor_ptr,
                                act_name,
                                egr::Controller::Instance().GetExpectedPlace(),
                                persistable,
                                /* stop_gradient */ -1,
                                dtype,
                                dims,
-                               var_type);
+                               var_type,
+                               process_mesh_ptr,
+                               placements_ptr);
 
         return 0;
       } else {
@@ -1025,12 +1061,12 @@ int TensorInit(PyObject* self, PyObject* args, PyObject* kwargs) {
           py_tensor_ptr, kws_map, args, flag_kwargs, args_num);
       return 0;
     } else if (PyObject_TypeCheck(arg0_ptr, p_tensor_type)) {
-      VLOG(6) << "Calling case5's or case6's initializer.";
+      VLOG(6) << "Calling case5's or case6's or case7's initializer.";
       AutoInitTensorByTensor(
           py_tensor_ptr, kws_map, args, flag_kwargs, args_num);
       return 0;
     } else if (PyObject_TypeCheck(arg0_ptr, g_framework_tensor_pytype)) {
-      VLOG(6) << "Calling case7's initializer.";
+      VLOG(6) << "Calling case8's initializer.";
       AutoInitTensorByTensor(py_tensor_ptr,
                              kws_map,
                              args,
@@ -1137,6 +1173,56 @@ int TensorInit(PyObject* self, PyObject* args, PyObject* kwargs) {
           "Please check your code and make sure the first position args is "
           "PyArray."));
     }
+  } else if (args_num == (Py_ssize_t)7) {
+    if (!flag_kwargs) {
+      PyObject* arg0_ptr = PyTuple_GET_ITEM(args, 0);
+      if (PyObject_TypeCheck(arg0_ptr, g_vartype_pytype)) {
+        VLOG(6) << "Calling case2's initializer.";
+        paddle::framework::proto::VarType::Type dtype =
+            CastPyArg2ProtoType(PyTuple_GET_ITEM(args, 0), 0);
+        std::vector<int> dims =
+            CastPyArg2VectorOfInt(PyTuple_GET_ITEM(args, 1), 1);
+        std::string act_name = "";
+        PyObject* name_obj = PyTuple_GET_ITEM(args, 2);
+        if (name_obj == Py_None) {
+          act_name = egr::Controller::Instance().GenerateUniqueName(
+              "generated_tensor");
+        } else {
+          act_name = CastPyArg2AttrString(PyTuple_GET_ITEM(args, 2), 2);
+        }
+        paddle::framework::proto::VarType::Type var_type =
+            CastPyArg2ProtoType(PyTuple_GET_ITEM(args, 3), 3);
+        bool persistable = CastPyArg2AttrBoolean(PyTuple_GET_ITEM(args, 4), 4);
+        ProcessMesh process_mesh =
+            CastPyArg2ProcessMesh(PyTuple_GET_ITEM(args, 5), 5);
+        Placements placements =
+            CastPyArg2VectorOfPlacement(PyTuple_GET_ITEM(args, 6), 6);
+        EmptyTensorInitializer(py_tensor_ptr,
+                               act_name,
+                               egr::Controller::Instance().GetExpectedPlace(),
+                               persistable,
+                               -1,
+                               dtype,
+                               dims,
+                               var_type,
+                               &process_mesh,
+                               &placements);
+        return 0;
+      } else {
+        PADDLE_THROW(platform::errors::InvalidArgument(
+            "Incompatible constructor arguments, "
+            "there are only 7 position args,"
+            "but the first position args should be dtype. "
+            "Please check your code and make sure you call the existed "
+            "constructor."));
+      }
+    } else {
+      PADDLE_THROW(platform::errors::InvalidArgument(
+          "Incompatible constructor arguments, "
+          "there are 7 position args and remaining arguments are kwargs,"
+          "Please check your code and make sure you call the existed "
+          "constructor."));
+    }
   } else {
     PADDLE_THROW(platform::errors::Fatal(
         "Can't not find expected num of args, please check your call, and "
diff --git a/python/paddle/distributed/fleet/recompute/recompute.py b/python/paddle/distributed/fleet/recompute/recompute.py
@@ -295,6 +295,8 @@ def _recompute_without_reentrant(
         cur_device = paddle.get_device()
         if 'gpu:' in cur_device:
             fw_cuda_rng_state = paddle.get_cuda_rng_state()
+        elif 'cpu' in cur_device:
+            fw_cuda_rng_state = paddle.get_rng_state()
         elif 'xpu:' in cur_device:
             fw_cuda_rng_state = paddle.get_rng_state()
         elif (
@@ -346,13 +348,26 @@ def inner_pack(inner_x):
                     return
 
                 if inner_x.is_contiguous():
-                    tmp_tensor = core.eager.Tensor(
-                        inner_x.dtype,
-                        inner_x.shape,
-                        inner_x.name + "cpy",
-                        core.VarDesc.VarType.LOD_TENSOR,
-                        inner_x.persistable,
-                    )
+                    if inner_x.is_dist():
+                        # TODO(jeff41404): it seems better to use `tmp_tensor = core.eager.Tensor(inner_x)`,
+                        # but other errors will be triggered during the current period, and can be modified after resolution
+                        tmp_tensor = core.eager.Tensor(
+                            inner_x.dtype,
+                            inner_x.shape,
+                            inner_x.name + "cpy",
+                            core.VarDesc.VarType.LOD_TENSOR,
+                            inner_x.persistable,
+                            inner_x.process_mesh,
+                            inner_x.placements,
+                        )
+                    else:
+                        tmp_tensor = core.eager.Tensor(
+                            inner_x.dtype,
+                            inner_x.shape,
+                            inner_x.name + "cpy",
+                            core.VarDesc.VarType.LOD_TENSOR,
+                            inner_x.persistable,
+                        )
                     inner_x._share_buffer_to(tmp_tensor)
                     storage[holder_list[unpack_counter - 1]()] = tmp_tensor
                 else:
diff --git a/test/auto_parallel/semi_auto_parallel_simple_net.py b/test/auto_parallel/semi_auto_parallel_simple_net.py
@@ -40,6 +40,7 @@ def __init__(
         self,
         param_prefix="",
         is_recompute=False,
+        recompute_use_reentrant=True,
         is_pp=False,
         pp_reshard_dist_attr=None,
     ):
@@ -50,6 +51,7 @@ def __init__(
 
         self.is_pp = is_pp
         self.is_recompute = is_recompute
+        self.recompute_use_reentrant = recompute_use_reentrant
         self.pp_reshard_dist_attr = pp_reshard_dist_attr
         self.linear_0 = nn.Linear(
             IMAGE_SIZE, IMAGE_SIZE, weight_attr_0, bias_attr=False
@@ -70,7 +72,10 @@ def _inner_forward_fn(self, x):
 
     def forward(self, x):
         if self.is_recompute:
-            return recompute(self._inner_forward_fn, x)
+            if self.recompute_use_reentrant:
+                return recompute(self._inner_forward_fn, x)
+            else:
+                return recompute(self._inner_forward_fn, x, use_reentrant=False)
         else:
             return self._inner_forward_fn(x)
 
diff --git a/test/auto_parallel/semi_auto_parallel_simple_net_recompute.py b/test/auto_parallel/semi_auto_parallel_simple_net_recompute.py
@@ -81,6 +81,25 @@ def test_dp_demo_net(self):
             self.check_tensor_eq(param, param_base, rtol=1e-4)
             self.check_tensor_eq(param.grad, param_base.grad)
 
+    def test_dp_demo_net_use_reentrant_false(self):
+        self.set_random_seed(self._seed)
+        (
+            self.dp_loss,
+            self.dp_parameters,
+        ) = self.run_dynamic_recompute(
+            DemoNet(
+                "recompute_use_reentrant_false_dp_demo",
+                is_recompute=True,
+                recompute_use_reentrant=False,
+            ),
+            shard_input=True,
+        )
+        self.check_tensor_eq(self.dp_loss, self.base_loss)
+        self.check_tensor_eq(self.dp_loss, self.base_loss)
+        for param, param_base in zip(self.dp_parameters, self.base_parameters):
+            self.check_tensor_eq(param, param_base, rtol=1e-4)
+            self.check_tensor_eq(param.grad, param_base.grad)
+
     def test_mp_demo_net(self):
         self.set_random_seed(self._seed)
         mp_layer = dist.shard_layer(
@@ -98,9 +117,32 @@ def test_mp_demo_net(self):
             self.check_tensor_eq(param, param_base)
             self.check_tensor_eq(param.grad, param_base.grad)
 
+    def test_mp_demo_net_use_reentrant_false(self):
+        self.set_random_seed(self._seed)
+        mp_layer = dist.shard_layer(
+            DemoNet(
+                "recompute_use_reentrant_false_mp_demo",
+                is_recompute=True,
+                recompute_use_reentrant=False,
+            ),
+            self._mesh,
+            self.shard_fn,
+        )
+        (
+            self.mp_loss,
+            self.mp_parameters,
+        ) = self.run_dynamic_recompute(mp_layer)
+
+        self.check_tensor_eq(self.mp_loss, self.base_loss)
+        for param, param_base in zip(self.mp_parameters, self.base_parameters):
+            self.check_tensor_eq(param, param_base)
+            self.check_tensor_eq(param.grad, param_base.grad)
+
     def run_test_case(self):
         self.test_dp_demo_net()
+        self.test_dp_demo_net_use_reentrant_false()
         self.test_mp_demo_net()
+        self.test_mp_demo_net_use_reentrant_false()
 
 
 if __name__ == '__main__':
