PaddlePaddle
diff --git a/‎paddle/fluid/eager/auto_code_generator/generator/eager_gen.py
Copy file name to clipboard
+23Lines changed: 23 additions & 0 deletions b/‎paddle/fluid/eager/auto_code_generator/generator/eager_gen.py
Copy file name to clipboard
+23Lines changed: 23 additions & 0 deletions
diff --git a/‎paddle/fluid/eager/grad_node_info.cc
Copy file name to clipboard
+6-3Lines changed: 6 additions & 3 deletions b/‎paddle/fluid/eager/grad_node_info.cc
Copy file name to clipboard
+6-3Lines changed: 6 additions & 3 deletions
diff --git a/‎paddle/fluid/eager/grad_node_info.h
Copy file name to clipboard
+25Lines changed: 25 additions & 0 deletions b/‎paddle/fluid/eager/grad_node_info.h
Copy file name to clipboard
+25Lines changed: 25 additions & 0 deletions
diff --git a/‎paddle/fluid/eager/grad_tensor_holder.cc
Copy file name to clipboard
+8-9Lines changed: 8 additions & 9 deletions b/‎paddle/fluid/eager/grad_tensor_holder.cc
Copy file name to clipboard
+8-9Lines changed: 8 additions & 9 deletions
diff --git a/‎paddle/fluid/eager/utils.h
Copy file name to clipboard
+93-7Lines changed: 93 additions & 7 deletions b/‎paddle/fluid/eager/utils.h
Copy file name to clipboard
+93-7Lines changed: 93 additions & 7 deletions
diff --git a/‎paddle/phi/api/lib/api_gen_utils.cc
Copy file name to clipboard
+5Lines changed: 5 additions & 0 deletions b/‎paddle/phi/api/lib/api_gen_utils.cc
Copy file name to clipboard
+5Lines changed: 5 additions & 0 deletions
diff --git a/‎paddle/phi/api/lib/api_gen_utils.h
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions b/‎paddle/phi/api/lib/api_gen_utils.h
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions
@@ -210,6 +210,8 @@ class {} : public egr::GradNodeBase {{
   // Prepare Grad function call
 {}
   // Runtime check if we need next grad
+{}
+  // Set DistAttr of Out Tensor for semi-auto parallel
 {}
   // Inplace Check
 {}
@@ -529,6 +531,12 @@ class {} : public egr::GradNodeBase {{
   if( !{}.empty() ) {}_optional = paddle::make_optional<std::vector<paddle::Tensor>>({});
 """
 
+SET_GRAD_OUT_DIST_ATTR_TEMPLATE = """
+  if (IsRunAutoParallel()) {{
+    egr::EagerUtils::SetGradOutputDistAttr(out_metas, {}, {});
+  }}
+"""
+
 CHECK_BACKWARD_INPLACE_TEMPLATE = """
   bool can_be_inplaced = false;
   if ({}.initialized()) {{
@@ -2181,6 +2189,8 @@ def GenerateNodeDefinition(
         )
         grad_api_args = ["" for i in range(grad_api_args_len)]
         get_grad_in_args_list = []
+        grad_api_out_args_list = []
+        fwd_positions_list = []
 
         # Fill Grad Ins with Zero
         fill_zero_str = ""
@@ -2388,6 +2398,8 @@ def GenerateNodeDefinition(
                     out_assign_str += f"{indent}*api_output_{out_index} = std::get<{out_index}>(api_output);\n"
             else:
                 grad_api_args.append(f"api_output_{out_index}")
+                grad_api_out_args_list.append(f"api_output_{out_index}")
+                fwd_positions_list.append(f"{fwd_position}")
             if inplace_grad_input_str in optional_inplace_var_name:
                 optional_inplace_str = "VLOG(6) << \"No Inplace should happend for wrappered input: {inplace_grad_input_str}\";"
             else:
@@ -2433,6 +2445,16 @@ def GenerateNodeDefinition(
         composite_grad_api_args_str = ", ".join(grad_api_args)
         composite_template_name = "<paddle::Tensor>"
 
+        # Set DistAttr Func Construct
+        set_out_dist_attr_str = ""
+        if not is_invoke_forward_api:
+            fwd_positions_str = "{" + ", ".join(fwd_positions_list) + "}"
+            grad_api_out_args_str = ", ".join(grad_api_out_args_list)
+            set_out_dist_attr_str = SET_GRAD_OUT_DIST_ATTR_TEMPLATE.format(
+                fwd_positions_str,
+                grad_api_out_args_str,
+            )
+
         if is_invoke_forward_api:
             autograd_api_out = "auto"
             if (
@@ -2600,6 +2622,7 @@ def GenerateNodeDefinition(
             get_grad_in_args_str,
             grad_function_prepare_str,
             compute_require_next_grad_str,
+            set_out_dist_attr_str,
             inplace_check_str,
             inplace_for_grad_outs_str,
             self.backward_api_name,
 
@@ -265,9 +265,9 @@ void GradNodeBase::SetGradOutMeta(const paddle::Tensor& fwd_in,
       meta.SetTensorMeta(dense_tensor->meta());
       meta.SetPlace(fwd_in.place());
     } else if (phi::distributed::DistTensor::classof(fwd_in.impl().get())) {
-      const phi::DenseTensor& dense_tensor =
-          static_cast<phi::distributed::DistTensor*>(fwd_in.impl().get())
-              ->value();
+      const phi::distributed::DistTensor* dist_tensor =
+          static_cast<phi::distributed::DistTensor*>(fwd_in.impl().get());
+      const phi::DenseTensor& dense_tensor = dist_tensor->value();
       PADDLE_ENFORCE_NE(
           dense_tensor.meta().dtype,
           phi::DataType::UNDEFINED,
@@ -276,6 +276,9 @@ void GradNodeBase::SetGradOutMeta(const paddle::Tensor& fwd_in,
                                           "which is illegal."));
       meta.SetTensorMeta(dense_tensor.meta());
       meta.SetPlace(fwd_in.place());
+      // Set DistAttr
+      meta.SetDistAttr(dist_tensor->dist_attr());
+      SetIsRunAutoParallel(true);
     } else {
       VLOG(7)
           << "Unable to initialize the DenseTensorMeta of GradSlotMeta with "
 
@@ -20,6 +20,7 @@
 #include "paddle/fluid/eager/eager_tensor.h"
 #include "paddle/fluid/eager/hooks.h"
 #include "paddle/phi/api/all.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
 
 namespace egr {
 /**
@@ -158,11 +159,23 @@ class GradSlotMeta {
   Edge& GetMutableEdge() { return adj_edge_; }
   const Edge& GetEdge() const { return adj_edge_; }
 
+  const phi::distributed::TensorDistAttr& DistAttr() const {
+    return dist_attr_;
+  }
+
+  void SetDistAttr(const phi::distributed::TensorDistAttr& dist_attr) {
+    dist_attr_ = dist_attr;
+  }
+
  private:
   bool stop_gradient_{false};
   phi::Place place_;
   std::shared_ptr<phi::DenseTensorMeta> meta_ = nullptr;
   Edge adj_edge_;
+  // For dygraph semi-auto parallel
+  // Save the dist attr of the forward input Tensor for proper resharding
+  // operation when compute the input Tensor's gradient
+  phi::distributed::TensorDistAttr dist_attr_;
 };
 
 class GradNodeBase {
@@ -298,6 +311,14 @@ class GradNodeBase {
 
   std::string GetForwardTrace() { return forward_trace_; }
 
+  /**
+   * The following interfaces are designed for auto parallel
+   * **/
+  bool IsRunAutoParallel() const { return is_run_auto_parallel_; }
+  void SetIsRunAutoParallel(bool is_run_auto_parallel) {
+    is_run_auto_parallel_ = is_run_auto_parallel;
+  }
+
  private:
   // bwd_out_meta_ is used to record Grad output info for backward
   paddle::small_vector<std::vector<GradSlotMeta>, kSlotSmallVectorSize>
@@ -325,6 +346,10 @@ class GradNodeBase {
   bool is_tensor_wrappers_cleared_ = false;
   // The trace of forward function
   std::string forward_trace_ = "";
+
+  // With this flag, short-circuit the backward traversal of Tensor and
+  // set the DistAttr to reduce the impact on scheduling performance
+  bool is_run_auto_parallel_{false};
 };
 
 }  // namespace egr
@@ -86,16 +86,15 @@ void GradTensorHolder::CopyValueFromTensor(size_t slot_id,
         buffer_[slot_id][rank] =
             paddle::experimental::sparse::full_like(t, 1, t.dtype());
       } else if (t.is_dist_tensor()) {
-        VLOG(6) << "Create a new dist tensor.";
-        // TODO(chenweihang): we need a shard_tensor API in C++
-        // TODO(chenweihang): replace by valid dist_attr later
-        auto temp =
+        auto init_grad =
             paddle::experimental::full(t.shape(), 1, t.dtype(), t.place());
-        auto dense_temp = static_cast<phi::DenseTensor*>(temp.impl().get());
-        auto dist_tensor = std::make_shared<phi::distributed::DistTensor>(
-            *dense_temp, phi::distributed::TensorDistAttr());
-        temp.set_impl(dist_tensor);
-        buffer_[slot_id][rank] = temp;
+        auto global_dense_t =
+            static_cast<phi::DenseTensor*>(init_grad.impl().get());
+        auto dist_t =
+            static_cast<phi::distributed::DistTensor*>(t.impl().get());
+        init_grad.set_impl(std::make_shared<phi::distributed::DistTensor>(
+            *global_dense_t, dist_t->dist_attr()));
+        buffer_[slot_id][rank] = init_grad;
       } else {
         PADDLE_THROW(paddle::platform::errors::Fatal(
             "Only Support DENSE_TENSOR, SPARSE_COO_TENSOR, SPARSE_CSR_TENSOR "
 
@@ -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/core/distributed/auto_parallel/dist_tensor.h"
 
 namespace egr {
 
@@ -33,10 +34,14 @@ template <typename ElementType>
 class IterHelper {
   virtual void visit(ElementType element) = 0;
 
-  void visit(std::vector<ElementType>* elements) {
+  virtual void visit(std::vector<ElementType>* elements) {
     for (auto element : *elements) visit(element);
   }
 
+  virtual void visit(const std::vector<ElementType>& elements) {
+    for (auto element : elements) visit(element);
+  }
+
   template <typename... Args>
   void apply() {}
 
@@ -82,6 +87,58 @@ class PassStopGradientIter : public IterHelper<AutogradMeta*> {
   bool stop_gradient_ = true;
 };
 
+class SetGradOutputDistAttrIter : public IterHelper<paddle::Tensor*> {
+ public:
+  explicit SetGradOutputDistAttrIter(
+      const paddle::small_vector<std::vector<GradSlotMeta>,
+                                 kSlotSmallVectorSize>& out_meta,
+      const paddle::small_vector<size_t, kSlotSmallVectorSize>& out_indexes)
+      : out_meta_(out_meta), out_indexes_{out_indexes} {}
+
+ private:
+  void visit_element(paddle::Tensor* element, const GradSlotMeta& meta) {
+    if (element == nullptr) {
+      return;
+    }
+    if (meta.DistAttr().empty()) {
+      return;
+    }
+    if (element->defined()) {
+      if (element->is_dist_tensor()) {
+        PADDLE_THROW(phi::errors::Unimplemented(
+            "Unsupport set defined dist tensor now."));
+      } else {
+        // Only deal with dist tensor here
+        return;
+      }
+    } else {
+      element->set_impl(std::make_shared<phi::distributed::DistTensor>(
+          phi::DDim(), meta.DistAttr()));
+    }
+  }
+  void visit(paddle::Tensor* element) override {
+    if (!out_meta_[out_indexes_[cur_pos_]].empty()) {
+      visit_element(element, out_meta_[out_indexes_[cur_pos_]][0]);
+    }
+    cur_pos_++;
+  }
+
+  void visit(const std::vector<paddle::Tensor*>& elements) override {
+    if (!out_meta_[out_indexes_[cur_pos_]].empty()) {
+      for (size_t i = 0; i < elements.size(); ++i) {
+        visit_element(elements.at(i), out_meta_[out_indexes_[cur_pos_]][i]);
+      }
+    }
+    cur_pos_++;
+  }
+
+  const paddle::small_vector<std::vector<GradSlotMeta>, kSlotSmallVectorSize>&
+      out_meta_;
+  const paddle::small_vector<size_t, kSlotSmallVectorSize>& out_indexes_;
+
+  int cur_pos_{0};
+};
+
 class EagerUtils {
  public:
   /**
@@ -223,6 +280,20 @@ class EagerUtils {
   static void FillZeroForEmptyGradInput(
       std::vector<paddle::Tensor>* in_grads,
       const std::vector<GradSlotMeta>& grad_in_metas);
+
+  /**
+   * Set DistAttr
+   */
+  template <typename... Args>
+  static void SetGradOutputDistAttr(
+      const paddle::small_vector<std::vector<GradSlotMeta>,
+                                 kSlotSmallVectorSize>& out_metas,
+      const paddle::small_vector<size_t, kSlotSmallVectorSize>& out_indexes,
+      Args&&... args) {
+    SetGradOutputDistAttrIter(out_metas, out_indexes)
+        .apply(std::forward<Args>(args)...);
+  }
+
   /**
    * Print Input Output (level 0 means least info, level 2 means most info)
    * **/
@@ -234,15 +305,30 @@ class EagerUtils {
       tensor_name_str = t.name();
     }
     const char* TENSOR_INFO_TEMPLATE =
-        "Type: %s, Dtype: %s, Place: %s, Shape: %s";
+        "Type: %s, Dtype: %s, Place: %s, Shape: %s, DistAttr: %s";
     std::string tensor_info_str = "";
     if (t.defined()) {
       if (t.initialized()) {
-        tensor_info_str += paddle::string::Sprintf(TENSOR_INFO_TEMPLATE,
-                                                   t.impl()->type_info().name(),
-                                                   t.dtype(),
-                                                   t.place().DebugString(),
-                                                   t.dims());
+        if (t.is_dist_tensor()) {
+          auto dist_t =
+              std::static_pointer_cast<phi::distributed::DistTensor>(t.impl());
+          tensor_info_str += paddle::string::Sprintf(
+              TENSOR_INFO_TEMPLATE,
+              t.impl()->type_info().name(),
+              t.dtype(),
+              t.place().DebugString(),
+              paddle::string::Sprintf(
+                  "%s, Local Shape: %s", t.dims(), dist_t->local_dims()),
+              dist_t->dist_attr());
+        } else {
+          tensor_info_str +=
+              paddle::string::Sprintf(TENSOR_INFO_TEMPLATE,
+                                      t.impl()->type_info().name(),
+                                      t.dtype(),
+                                      t.place().DebugString(),
+                                      t.dims(),
+                                      "Unknown");
+        }
       } else {
         tensor_info_str += paddle::string::Sprintf(TENSOR_INFO_TEMPLATE,
                                                    t.impl()->type_info().name(),
 
@@ -550,6 +550,11 @@ phi::distributed::DistTensor* SetKernelDistOutput(
   return nullptr;
 }
 
+std::shared_ptr<phi::distributed::DistTensor> CreateKernelDistOutput(
+    const phi::distributed::TensorDistAttr& dist_attr) {
+  return std::make_shared<phi::distributed::DistTensor>(phi::DDim(), dist_attr);
+}
+
 std::vector<phi::distributed::DistTensor*> SetKernelDistOutput(
     std::vector<Tensor*> out) {
   std::vector<phi::distributed::DistTensor*> result;
 
@@ -144,6 +144,9 @@ phi::distributed::DistTensor* SetKernelDistOutput(
     const phi::distributed::TensorDistAttr& dist_attr =
         phi::distributed::TensorDistAttr());
 
+std::shared_ptr<phi::distributed::DistTensor> CreateKernelDistOutput(
+    const phi::distributed::TensorDistAttr& dist_attr);
+
 std::vector<phi::distributed::DistTensor*> SetKernelDistOutput(
     std::vector<Tensor*> out);