convert disttensor for pylayer (#59148)

wanghuancoder · web-flow · commit 831137c89f28 · 2023-11-20T20:16:43.000-06:00
diff --git a/paddle/fluid/pybind/eager_py_layer.cc b/paddle/fluid/pybind/eager_py_layer.cc
@@ -172,6 +172,54 @@ PyObject* pylayer_method_apply(PyObject* cls,
   ctx->forward_input_tensor_is_duplicable.clear();
   ctx->forward_input_tensor_is_duplicable.reserve(inputs_size);
   std::set<phi::TensorBase*> input_tensorbases;
+
+  const phi::distributed::ProcessMesh* mesh = nullptr;
+  for (size_t i = 0; i < inputs_size; i++) {
+    PyObject* obj = nullptr;
+    if (i >= args_size) {
+      obj = PyList_GetItem(kwargs_value_list, i - args_size);  // NOLINT
+    } else {
+      obj = PyTuple_GET_ITEM(args, i);
+    }
+    if (PyCheckTensor(obj)) {
+      paddle::Tensor& tensor = reinterpret_cast<TensorObject*>(obj)->tensor;
+      if (tensor.defined() && tensor.is_dist_tensor()) {
+        mesh = &(std::dynamic_pointer_cast<phi::distributed::DistTensor>(
+                     tensor.impl())
+                     ->dist_attr()
+                     .process_mesh());
+      }
+    } else if (PyList_Check(obj)) {
+      Py_ssize_t len = PyList_Size(obj);
+      for (Py_ssize_t j = 0; j < len; j++) {
+        PyObject* o = PyList_GetItem(obj, j);
+        if (PyCheckTensor(o)) {
+          paddle::Tensor& tensor = reinterpret_cast<TensorObject*>(o)->tensor;
+          if (tensor.defined() && tensor.is_dist_tensor()) {
+            mesh = &(std::dynamic_pointer_cast<phi::distributed::DistTensor>(
+                         tensor.impl())
+                         ->dist_attr()
+                         .process_mesh());
+          }
+        }
+      }
+    } else if (PyTuple_Check(obj)) {
+      Py_ssize_t len = PyTuple_Size(obj);
+      for (Py_ssize_t j = 0; j < len; j++) {
+        PyObject* o = PyTuple_GetItem(obj, j);
+        if (PyCheckTensor(o)) {
+          paddle::Tensor& tensor = reinterpret_cast<TensorObject*>(o)->tensor;
+          if (tensor.defined() && tensor.is_dist_tensor()) {
+            mesh = &(std::dynamic_pointer_cast<phi::distributed::DistTensor>(
+                         tensor.impl())
+                         ->dist_attr()
+                         .process_mesh());
+          }
+        }
+      }
+    }
+  }
+
   for (size_t i = 0; i < inputs_size; i++) {
     PyObject* obj = nullptr;
     if (i >= args_size) {
@@ -180,6 +228,10 @@ PyObject* pylayer_method_apply(PyObject* cls,
       obj = PyTuple_GET_ITEM(args, i);
     }
     if (PyCheckTensor(obj)) {
+      if (mesh) {
+        ConvertToDistTensor(&(reinterpret_cast<TensorObject*>(obj)->tensor),
+                            mesh);
+      }
       input_tensorbases.insert(
           reinterpret_cast<TensorObject*>(obj)->tensor.impl().get());
       auto autograd_meta = egr::EagerUtils::nullable_autograd_meta(
@@ -199,6 +251,10 @@ PyObject* pylayer_method_apply(PyObject* cls,
       for (Py_ssize_t j = 0; j < len; j++) {
         PyObject* o = PyList_GetItem(obj, j);
         if (PyCheckTensor(o)) {
+          if (mesh) {
+            ConvertToDistTensor(&(reinterpret_cast<TensorObject*>(o)->tensor),
+                                mesh);
+          }
           input_tensorbases.insert(
               reinterpret_cast<TensorObject*>(o)->tensor.impl().get());
           tensors.push_back(&(reinterpret_cast<TensorObject*>(o)->tensor));
@@ -222,6 +278,10 @@ PyObject* pylayer_method_apply(PyObject* cls,
       for (Py_ssize_t j = 0; j < len; j++) {
         PyObject* o = PyTuple_GetItem(obj, j);
         if (PyCheckTensor(o)) {
+          if (mesh) {
+            ConvertToDistTensor(&(reinterpret_cast<TensorObject*>(o)->tensor),
+                                mesh);
+          }
           input_tensorbases.insert(
               reinterpret_cast<TensorObject*>(o)->tensor.impl().get());
           tensors.push_back(&(reinterpret_cast<TensorObject*>(o)->tensor));
diff --git a/paddle/fluid/pybind/eager_utils.h b/paddle/fluid/pybind/eager_utils.h
@@ -465,5 +465,6 @@ void ConvertAllInputsToDistTensor(const phi::distributed::ProcessMesh* mesh,
   DistTensorConverter(mesh).apply(&args...);
 }
 
+void ConvertToDistTensor(Tensor* x, const phi::distributed::ProcessMesh* mesh);
 }  // namespace pybind
 }  // namespace paddle

Original file line number	Diff line number	Diff line change
`@@ -465,5 +465,6 @@ void ConvertAllInputsToDistTensor(const phi::distributed::ProcessMesh* mesh,`
`465`	`465`	`DistTensorConverter(mesh).apply(&args...);`
`466`	`466`	`}`
`467`	`467`
	`468`	`+void ConvertToDistTensor(Tensor* x, const phi::distributed::ProcessMesh* mesh);`
`468`	`469`	`} // namespace pybind`
`469`	`470`	`} // namespace paddle`