PaddlePaddle
diff --git a/‎paddle/fluid/eager/api/manual/eager_manual/nodes/multiply_node.cc
Copy file name to clipboard
+14Lines changed: 14 additions & 0 deletions b/‎paddle/fluid/eager/api/manual/eager_manual/nodes/multiply_node.cc
Copy file name to clipboard
+14Lines changed: 14 additions & 0 deletions
diff --git a/‎paddle/fluid/eager/auto_code_generator/generator/python_c_gen.py
Copy file name to clipboard
+4-3Lines changed: 4 additions & 3 deletions b/‎paddle/fluid/eager/auto_code_generator/generator/python_c_gen.py
Copy file name to clipboard
+4-3Lines changed: 4 additions & 3 deletions
diff --git a/‎paddle/fluid/eager/utils.cc
Copy file name to clipboard
+155Lines changed: 155 additions & 0 deletions b/‎paddle/fluid/eager/utils.cc
Copy file name to clipboard
+155Lines changed: 155 additions & 0 deletions
diff --git a/‎paddle/fluid/eager/utils.h
Copy file name to clipboard
+74Lines changed: 74 additions & 0 deletions b/‎paddle/fluid/eager/utils.h
Copy file name to clipboard
+74Lines changed: 74 additions & 0 deletions
diff --git a/‎paddle/fluid/pybind/eager_custom_python_api.h
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions b/‎paddle/fluid/pybind/eager_custom_python_api.h
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions
diff --git a/‎paddle/fluid/pybind/eager_functions.cc
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions b/‎paddle/fluid/pybind/eager_functions.cc
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions
diff --git a/‎paddle/fluid/pybind/eager_legacy_custom_python_api.h
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions b/‎paddle/fluid/pybind/eager_legacy_custom_python_api.h
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions
diff --git a/‎paddle/fluid/pybind/eager_math_op_patch.cc
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions b/‎paddle/fluid/pybind/eager_math_op_patch.cc
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions
diff --git a/‎paddle/fluid/pybind/eager_method.cc
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions b/‎paddle/fluid/pybind/eager_method.cc
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions
diff --git a/‎paddle/fluid/pybind/eager_properties.cc
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions b/‎paddle/fluid/pybind/eager_properties.cc
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions
@@ -31,6 +31,9 @@
 #include "paddle/phi/api/lib/api_custom_impl.h"
 #include "paddle/phi/core/platform/profiler/event_tracing.h"
 
+using egr::ConvertAllInputsToDistTensor;
+using egr::InputsContainDistTensor;
+
 COMMON_DECLARE_bool(check_nan_inf);
 
 paddle::small_vector<std::vector<paddle::Tensor>, egr::kSlotSmallVectorSize>
@@ -64,6 +67,14 @@ MultiplyGradNode::operator()(
   auto y = egr::EagerUtils::RecoverTensorWrapper(&this->y_);
   auto& grad_out = hooked_grads[0][0];
   auto& axis = this->axis_;
+
+  // Convert All Inputs to DistTensor if Necessary
+  const phi::distributed::ProcessMesh* mesh = nullptr;
+  bool inputs_contain_dist_tensor = InputsContainDistTensor(&mesh, grad_out);
+  if (inputs_contain_dist_tensor) {
+    ConvertAllInputsToDistTensor(mesh, x, y);
+  }
+
   // Prepare Grad function call
 
   const auto& out_metas = OutputMeta();
@@ -89,6 +100,9 @@ MultiplyGradNode::operator()(
   if (IsRunAutoParallel()) {
     egr::EagerUtils::SetGradOutputDistAttr(
         out_metas, {0, 1}, api_output_0, api_output_1);
+  } else if (inputs_contain_dist_tensor) {
+    egr::EagerUtils::SetGradOutputDistAttrByInput(x, api_output_0);
+    egr::EagerUtils::SetGradOutputDistAttrByInput(y, api_output_1);
   }
 
   // Inplace Check
 
@@ -82,15 +82,15 @@ def FindParsingFunctionFromAttributeType(atype):
 
 CONVERT_INPUT_TENSORS_TO_DIST_TENSOR_WITH_SINGLE_TENSOR_TEMPLATE = """
     const phi::distributed::ProcessMesh* mesh = nullptr;
-    if (InputsContainDistTensor(&mesh{input_names})) {{
-      ConvertAllInputsToDistTensor(mesh{input_single_tensor_names});
+    if (egr::InputsContainDistTensor(&mesh{input_names})) {{
+      egr::ConvertAllInputsToDistTensor(mesh{input_single_tensor_names});
       {optional_and_vector_convert_code}
     }}
 """
 
 CONVERT_INPUT_TENSORS_TO_DIST_TENSOR_WITHOUT_SINGLE_TENSOR_TEMPLATE = """
     const phi::distributed::ProcessMesh* mesh = nullptr;
-    if (InputsContainDistTensor(&mesh{input_names})) {{
+    if (egr::InputsContainDistTensor(&mesh{input_names})) {{
       {optional_and_vector_convert_code}
     }}
 """
@@ -193,6 +193,7 @@ def FindParsingFunctionFromAttributeType(atype):
 #include "paddle/fluid/pybind/op_function_common.h"
 #include "paddle/fluid/eager/api/generated/eager_generated/forwards/dygraph_functions.h"
 #include "paddle/fluid/eager/api/manual/eager_manual/dygraph_forward_api.h"
+#include "paddle/fluid/eager/utils.h"
 #include "paddle/fluid/pybind/eager_custom_python_api.h"
 #include "paddle/fluid/pybind/eager.h"
 #include "paddle/fluid/pybind/eager_op_function.h"
 
@@ -21,6 +21,7 @@
 
 #include "paddle/common/layout.h"
 #include "paddle/phi/api/all.h"
+#include "paddle/phi/api/lib/data_transform.h"
 #include "paddle/phi/core/compat/convert_utils.h"
 #include "paddle/phi/core/tensor_meta.h"
 
@@ -943,4 +944,158 @@ std::string EagerUtils::TensorStr(
     return "[ " + tensors_str + " ]";
   }
 }
+
+void EagerUtils::SetGradOutputDistAttrByInput(const paddle::Tensor& input,
+                                              paddle::Tensor* grad) {
+  if (!grad) {
+    VLOG(4) << "The grad element is nullptr when calling "
+               "SetGradOutputDistAttrByInput.";
+    return;
+  }
+  auto dist_attr =
+      std::static_pointer_cast<phi::distributed::DistTensor>(input.impl())
+          ->dist_attr();
+  grad->set_impl(
+      std::make_shared<phi::distributed::DistTensor>(phi::DDim(), dist_attr));
+  VLOG(4) << "The grad element is set DistTensor impl when calling "
+             "SetGradOutputDistAttrByInput.";
+}
+
+void DistTensorTypeParser::operator()(const paddle::Tensor& x) {
+  if (x.defined() && x.is_dist_tensor()) {
+    *mesh = &(std::dynamic_pointer_cast<phi::distributed::DistTensor>(x.impl())
+                  ->process_mesh());
+    result = true;
+  }
+}
+
+void DistTensorTypeParser::operator()(
+    const paddle::optional<paddle::Tensor>& x) {
+  if (x) {
+    if (x.get_ptr()->defined() && x.get_ptr()->is_dist_tensor()) {
+      *mesh = &(std::dynamic_pointer_cast<phi::distributed::DistTensor>(
+                    x.get_ptr()->impl())
+                    ->process_mesh());
+      result = true;
+    }
+  }
+}
+
+void DistTensorTypeParser::operator()(const std::vector<paddle::Tensor>& x) {
+  if (!x.empty()) {
+    for (auto& t : x) {
+      if (t.defined() && t.is_dist_tensor()) {
+        *mesh =
+            &(std::dynamic_pointer_cast<phi::distributed::DistTensor>(t.impl())
+                  ->process_mesh());
+        result = true;
+        break;
+      }
+    }
+  }
+}
+
+void DistTensorTypeParser::operator()(
+    const paddle::optional<std::vector<paddle::Tensor>>& x) {
+  if (x) {
+    if (!(x.get_ptr()->empty())) {
+      for (auto& t : *(x.get_ptr())) {
+        if (t.defined() && t.is_dist_tensor()) {
+          *mesh = &(
+              std::dynamic_pointer_cast<phi::distributed::DistTensor>(t.impl())
+                  ->process_mesh());
+          result = true;
+          break;
+        }
+      }
+    }
+  }
+}
+
+void DistTensorConverter::convert(paddle::Tensor* x) {
+  ConvertToDistTensor(x, mesh);
+}
+
+void DistTensorConverter::operator()(paddle::Tensor* x) {
+  DistTensorConverter::convert(x);
+}
+
+void DistTensorConverter::operator()(paddle::optional<paddle::Tensor>* x) {
+  if (*x) {
+    DistTensorConverter::convert(x->get_ptr());
+  }
+}
+
+void DistTensorConverter::operator()(std::vector<paddle::Tensor>* x) {
+  if (!x->empty()) {
+    for (auto& t : *x) {
+      DistTensorConverter::convert(&t);
+    }
+  }
+}
+
+void DistTensorConverter::operator()(
+    paddle::optional<std::vector<paddle::Tensor>>* x) {
+  if (*x) {
+    if (!(x->get_ptr()->empty())) {
+      for (auto& t : *(x->get_ptr())) {
+        if (!t.is_dist_tensor()) {
+          DistTensorConverter::convert(&t);
+        }
+      }
+    }
+  }
+}
+
+void ConvertToDistTensor(paddle::Tensor* x,
+                         const phi::distributed::ProcessMesh* mesh) {
+  if (!x->defined()) {
+    return;
+  }
+  if (x->is_dist_tensor()) {
+    auto dist_ptr =
+        std::dynamic_pointer_cast<phi::distributed::DistTensor>(x->impl());
+    if (!dist_ptr->skip_check_mesh() && x->dims().size() > 0) {
+      // NOTE(pkuzyc): In MoE expert parallelism, the mesh of the
+      // inputs and outputs of different experts are different, so
+      // skip checking mesh in the following two casees:
+      // 1. The ``skip_check_mesh_`` flag is true. The MoE-related apis
+      // sets this flag to indicate that the difference between tensor's
+      // mesh is allowed.
+      // 2. The tensor is a 0-D tensor. Specifically, in MoE expert
+      // parallelism, the learning rate's mesh is global, but expert
+      // weights' mesh is the subset of the global mesh, this is also
+      // allowed so skip checking the mesh of 0-D tensor.
+      PADDLE_ENFORCE_EQ(
+          std::dynamic_pointer_cast<phi::distributed::DistTensor>(x->impl())
+              ->process_mesh(),
+          *mesh,
+          common::errors::InvalidArgument(
+              "Input %s has different mesh. However all inputs should "
+              "have the same mesh.",
+              x->name()));
+    }
+    return;
+  } else {
+    PADDLE_ENFORCE_EQ(
+        phi::DenseTensor::classof(x->impl().get()),
+        true,
+        common::errors::InvalidArgument(
+            "Failed to convert input %s impl to phi::distributed::DistTensor "
+            "as it's not phi::DenseTensor.",
+            x->name()));
+    phi::distributed::Placements placements;
+    for (int64_t i = 0; i < mesh->ndim(); ++i) {
+      placements.emplace_back(std::make_shared<phi::distributed::Replicate>());
+    }
+
+    auto dense_t = std::static_pointer_cast<phi::DenseTensor>(x->impl());
+    // auto parallel in dygraph doesn't support strided kernel.
+    if (!dense_t->meta().is_contiguous()) {
+      *dense_t = paddle::experimental::Trans2Contiguous(*dense_t);
+    }
+    x->set_impl(std::make_shared<phi::distributed::DistTensor>(
+        dense_t, *mesh, placements));
+  }
+}
 }  // namespace egr
@@ -18,6 +18,7 @@
 #include "paddle/fluid/eager/eager_tensor.h"
 #include "paddle/fluid/eager/grad_node_info.h"
 #include "paddle/phi/api/all.h"
+#include "paddle/phi/api/lib/kernel_dispatch.h"
 #include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
 #include "paddle/utils/test_macros.h"
 
@@ -268,6 +269,12 @@ class TEST_API EagerUtils {
         .apply(std::forward<Args>(args)...);
   }
 
+  /**
+   * Set DistAttr By Input
+   */
+  static void SetGradOutputDistAttrByInput(const paddle::Tensor& input,
+                                           paddle::Tensor* grad);
+
   /**
    * Print Input Output (level 0 means least info, level 2 means most info)
    * **/
@@ -284,4 +291,71 @@ class TEST_API EagerUtils {
       const paddle::optional<std::vector<paddle::Tensor>>& tensors);
 };
 
+using paddle::experimental::detail::ArgsIterator;
+
+struct DistTensorTypeParser : ArgsIterator<DistTensorTypeParser> {
+  bool result = false;
+  const phi::distributed::ProcessMesh** mesh = nullptr;
+
+  explicit DistTensorTypeParser(const phi::distributed::ProcessMesh** m)
+      : mesh(m) {}
+
+  bool short_circuit() { return result; }
+
+  void operator()(const paddle::Tensor& x);
+  void operator()(const paddle::optional<paddle::Tensor>& x);
+  void operator()(const std::vector<paddle::Tensor>& x);
+  void operator()(const paddle::optional<std::vector<paddle::Tensor>>& x);
+
+  // skip other type args, these args don't used in kernel selection
+  template <typename T>
+  void operator()(const T& x) {
+    // do nothing
+  }
+};
+
+struct DistTensorConverter : ArgsIterator<DistTensorConverter> {
+  const phi::distributed::ProcessMesh* mesh = nullptr;
+
+  explicit DistTensorConverter(const phi::distributed::ProcessMesh* m) {
+    PADDLE_ENFORCE_NE(
+        m,
+        nullptr,
+        common::errors::InvalidArgument(
+            "Input mesh of DistTensorConverter() shouldn't be nullptr."));
+    mesh = m;
+  }
+
+  void convert(paddle::Tensor* x);
+  void operator()(paddle::Tensor* x);
+  void operator()(paddle::optional<paddle::Tensor>* x);
+  void operator()(std::vector<paddle::Tensor>* x);
+  void operator()(paddle::optional<std::vector<paddle::Tensor>>* x);
+
+  // skip other type args, these args don't used in kernel selection
+  template <typename T>
+  void operator()(const T& x) {
+    // do nothing
+  }
+};
+
+template <typename... Args>
+bool InputsContainDistTensor(const phi::distributed::ProcessMesh** mesh,
+                             const Args&... args) {
+  return DistTensorTypeParser(mesh).apply(args...).result;
+}
+
+template <typename... Args>
+void ConvertAllInputsToDistTensor(const phi::distributed::ProcessMesh* mesh,
+                                  Args&... args) {
+  PADDLE_ENFORCE_NE(
+      mesh,
+      nullptr,
+      common::errors::InvalidArgument("Input mesh should not be nullptr."));
+  DistTensorConverter(mesh).apply(&args...);
+}
+
+void ConvertToDistTensor(paddle::Tensor* x,
+                         const phi::distributed::ProcessMesh* mesh);
+
 }  // namespace egr
@@ -15,8 +15,12 @@
 
 #include <iostream>
 
+#include "paddle/fluid/eager/utils.h"
 #include "paddle/phi/core/enforce.h"
 
+using egr::ConvertAllInputsToDistTensor;
+using egr::InputsContainDistTensor;
+
 namespace paddle {
 namespace pybind {
 
 
@@ -75,6 +75,9 @@ typedef SSIZE_T ssize_t;
 
 COMMON_DECLARE_string(tensor_operants_mode);
 
+using egr::ConvertAllInputsToDistTensor;
+using egr::InputsContainDistTensor;
+
 namespace paddle::pybind {
 
 namespace py = ::pybind11;
 
@@ -16,8 +16,12 @@
 #include <iostream>
 
 #include "paddle/fluid/eager/to_static/run_program_op_func.h"
+#include "paddle/fluid/eager/utils.h"
 #include "paddle/phi/core/enforce.h"
 
+using egr::ConvertAllInputsToDistTensor;
+using egr::InputsContainDistTensor;
+
 namespace paddle {
 namespace pybind {
 
 
@@ -54,6 +54,9 @@ typedef SSIZE_T ssize_t;
 #include "paddle/phi/core/memory/allocation/mmap_allocator.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
 
+using egr::ConvertAllInputsToDistTensor;
+using egr::InputsContainDistTensor;
+
 namespace paddle::pybind {
 
 static bool PyCheckInteger(PyObject* obj) {
 
@@ -72,6 +72,9 @@ typedef SSIZE_T ssize_t;
 COMMON_DECLARE_bool(set_to_1d);
 COMMON_DECLARE_bool(use_stride_kernel);
 
+using egr::ConvertAllInputsToDistTensor;
+using egr::InputsContainDistTensor;
+
 namespace paddle::pybind {
 
 extern void InitTensorWithNumpyValue(TensorObject* self,
 
@@ -33,6 +33,9 @@ limitations under the License. */
 #include "paddle/phi/core/memory/allocation/allocator.h"
 #include "paddle/phi/core/memory/memcpy.h"
 
+using egr::ConvertAllInputsToDistTensor;
+using egr::InputsContainDistTensor;
+
 #pragma GCC diagnostic ignored "-Wwrite-strings"
 
 COMMON_DECLARE_bool(enable_pir_api);