PaddlePaddle
diff --git a/‎paddle/fluid/eager/auto_code_generator/generator/eager_gen.py
Copy file name to clipboard
+3-9Lines changed: 3 additions & 9 deletions b/‎paddle/fluid/eager/auto_code_generator/generator/eager_gen.py
Copy file name to clipboard
+3-9Lines changed: 3 additions & 9 deletions
diff --git a/‎paddle/fluid/eager/grad_node_info.cc
Copy file name to clipboard
+97-4Lines changed: 97 additions & 4 deletions b/‎paddle/fluid/eager/grad_node_info.cc
Copy file name to clipboard
+97-4Lines changed: 97 additions & 4 deletions
diff --git a/‎paddle/fluid/eager/grad_node_info.h
Copy file name to clipboard
+14Lines changed: 14 additions & 0 deletions b/‎paddle/fluid/eager/grad_node_info.h
Copy file name to clipboard
+14Lines changed: 14 additions & 0 deletions
diff --git a/‎paddle/fluid/eager/utils.cc
Copy file name to clipboard
+83-21Lines changed: 83 additions & 21 deletions b/‎paddle/fluid/eager/utils.cc
Copy file name to clipboard
+83-21Lines changed: 83 additions & 21 deletions
@@ -581,21 +581,15 @@ class {} : public egr::GradNodeBase {{
 """
 
 FILL_ZERO_GRAD_TEMPLATE_BACKWARD = """
-  if (!IsRunAutoParallel()) {{
-    egr::EagerUtils::FillZeroForEmptyGradInput(&grads[{fwd_position}], input_metas[{fwd_position}]);
-  }}
+  egr::EagerUtils::FillZeroForEmptyGradInput(&grads[{fwd_position}], input_metas[{fwd_position}]);
 """
 
 FILL_ZERO_PLAIN_GRAD_TEMPLATE_BACKWARD = """
-  if (!IsRunAutoParallel()) {{
-    egr::EagerUtils::FillZeroForEmptyGradInput(&grads[{fwd_position}][0], input_metas[{fwd_position}][0]);
-  }}
+  egr::EagerUtils::FillZeroForEmptyGradInput(&grads[{fwd_position}][0], input_metas[{fwd_position}][0]);
 """
 
 FILL_ZERO_OPTIONAL_PLAIN_GRAD_TEMPLATE_BACKWARD = """
-  if (!IsRunAutoParallel()) {{
-      egr::EagerUtils::FillZeroForEmptyOptionalGradInput(&grads[{fwd_position}][0], input_metas[{fwd_position}][0]);
-  }}
+  egr::EagerUtils::FillZeroForEmptyOptionalGradInput(&grads[{fwd_position}][0], input_metas[{fwd_position}][0]);
 """
 
 inplace_optional_out_type_map = {
 
@@ -131,11 +131,16 @@ void GradNodeBase::SetGradInMeta(const paddle::Tensor& fwd_out,
         static_cast<phi::SparseCsrTensor*>(fwd_out.impl().get());
     dense_tensor = csr_tensor->mutable_non_zero_elements();
   } else if (phi::distributed::DistTensor::classof(fwd_out.impl().get())) {
-    // TODO(chenweihang): DistTensor contains global and local meta, here
-    // only set the local meta now, we should set global meta later
     dense_tensor =  // NOLINT
         &(static_cast<phi::distributed::DistTensor*>(fwd_out.impl().get())
               ->value());
+    meta.SetDistAttr(
+        static_cast<phi::distributed::DistTensor*>(fwd_out.impl().get())
+            ->dist_attr());
+    meta.SetDistTensorGlobalDims(
+        static_cast<phi::distributed::DistTensor*>(fwd_out.impl().get())
+            ->dims());
+    SetIsRunAutoParallel(true);
   } else {
     VLOG(7) << "Unable to initialize the DenseTensorMeta of GradSlotMeta with "
                "non-DenseTensor argument.";
@@ -221,6 +226,34 @@ void GradNodeBase::SetGradInMeta(const std::vector<paddle::Tensor>& fwd_out,
           dense_tensor->type() == phi::DataType::COMPLEX128) {
         need_complex_to_real_ = true;
       }
+    } else if (phi::distributed::DistTensor::classof(
+                   fwd_out_tensor.impl().get())) {
+      // Only Copy Meta
+      meta.SetDistAttr(static_cast<phi::distributed::DistTensor*>(
+                           fwd_out_tensor.impl().get())
+                           ->dist_attr());
+      meta.SetDistTensorGlobalDims(static_cast<phi::distributed::DistTensor*>(
+                                       fwd_out_tensor.impl().get())
+                                       ->dims());
+      SetIsRunAutoParallel(true);
+
+      auto dense_tensor = static_cast<phi::distributed::DistTensor*>(
+                              fwd_out_tensor.impl().get())
+                              ->value();
+
+      PADDLE_ENFORCE_NE(
+          dense_tensor.meta().dtype,
+          phi::DataType::UNDEFINED,
+          paddle::platform::errors::Fatal("Attempting to copy DenseTensorMeta "
+                                          "with phi::DataType::UNDEFINED,"
+                                          "which is illegal."));
+      meta.SetTensorMeta(dense_tensor.meta());
+      meta.SetPlace(fwd_out_tensor.place());
+
+      if (dense_tensor.type() == phi::DataType::COMPLEX64 ||
+          dense_tensor.type() == phi::DataType::COMPLEX128) {
+        need_complex_to_real_ = true;
+      }
     } else {
       VLOG(7) << "Unable to initialize the DenseTensorMeta of GradSlotMeta "
                  "with non-DenseTensor argument.";
@@ -298,6 +331,7 @@ void GradNodeBase::SetGradOutMeta(const paddle::Tensor& fwd_in,
           phi::errors::InvalidArgument(
               "The forward input DistTensor's dist attr is empty."));
       meta.SetDistAttr(dist_tensor->dist_attr());
+      meta.SetDistTensorGlobalDims(dist_tensor->dims());
       SetIsRunAutoParallel(true);
     } else {
       VLOG(7)
@@ -366,6 +400,26 @@ void GradNodeBase::SetGradOutMeta(const paddle::Tensor& fwd_in,
                                           "which is illegal."));
       meta.SetTensorMeta(dense_tensor->meta());
       meta.SetPlace(fwd_in.place());
+    } else if (phi::distributed::DistTensor::classof(fwd_in.impl().get())) {
+      // Only Copy Meta
+      meta.SetDistAttr(
+          static_cast<phi::distributed::DistTensor*>(fwd_in.impl().get())
+              ->dist_attr());
+      meta.SetDistTensorGlobalDims(
+          static_cast<phi::distributed::DistTensor*>(fwd_in.impl().get())
+              ->dims());
+      SetIsRunAutoParallel(true);
+      auto dense_tensor =
+          static_cast<phi::distributed::DistTensor*>(fwd_in.impl().get())
+              ->value();
+      PADDLE_ENFORCE_NE(
+          dense_tensor.meta().dtype,
+          phi::DataType::UNDEFINED,
+          paddle::platform::errors::Fatal("Attempting to copy DenseTensorMeta "
+                                          "with phi::DataType::UNDEFINED,"
+                                          "which is illegal."));
+      meta.SetTensorMeta(dense_tensor.meta());
+      meta.SetPlace(fwd_in.place());
     }
   } else {
     VLOG(7) << "Unable to initialize the DenseTensorMeta of GradSlotMeta with "
@@ -414,8 +468,6 @@ void GradNodeBase::SetGradOutMeta(const std::vector<paddle::Tensor>& fwd_in,
     // Record TensorMeta
     if (fwd_in_tensor.impl() && fwd_in_tensor.impl().get()) {
       if (phi::DenseTensor::classof(fwd_in_tensor.impl().get())) {
-        // TODO(chenweihang): DistTensor contains global and local meta, here
-        // only set the local meta now, we should set global meta later
         phi::DenseTensor* dense_tensor =
             static_cast<phi::DenseTensor*>(fwd_in_tensor.impl().get());
         PADDLE_ENFORCE_NE(dense_tensor->dtype(),
@@ -426,6 +478,26 @@ void GradNodeBase::SetGradOutMeta(const std::vector<paddle::Tensor>& fwd_in,
                               "which is illegal."));
         meta.SetTensorMeta(dense_tensor->meta());
         meta.SetPlace(fwd_in_tensor.place());
+      } else if (phi::distributed::DistTensor::classof(
+                     fwd_in_tensor.impl().get())) {
+        meta.SetDistAttr(static_cast<phi::distributed::DistTensor*>(
+                             fwd_in_tensor.impl().get())
+                             ->dist_attr());
+        meta.SetDistTensorGlobalDims(static_cast<phi::distributed::DistTensor*>(
+                                         fwd_in_tensor.impl().get())
+                                         ->dims());
+        SetIsRunAutoParallel(true);
+        auto dense_tensor = static_cast<phi::distributed::DistTensor*>(
+                                fwd_in_tensor.impl().get())
+                                ->value();
+        PADDLE_ENFORCE_NE(dense_tensor.dtype(),
+                          phi::DataType::UNDEFINED,
+                          paddle::platform::errors::Fatal(
+                              "Attempting to copy DenseTensorMeta "
+                              "with phi::DataType::UNDEFINED,"
+                              "which is illegal."));
+        meta.SetTensorMeta(dense_tensor.meta());
+        meta.SetPlace(fwd_in_tensor.place());
       }
     } else {
       VLOG(7)
@@ -487,6 +559,27 @@ void GradNodeBase::SetGradOutMeta(
                               "which is illegal."));
         meta.SetTensorMeta(dense_tensor->meta());
         meta.SetPlace(fwd_in_tensor.place());
+      } else if (phi::distributed::DistTensor::classof(
+                     fwd_in_tensor.impl().get())) {
+        // Only Copy Meta
+        meta.SetDistAttr(static_cast<phi::distributed::DistTensor*>(
+                             fwd_in_tensor.impl().get())
+                             ->dist_attr());
+        meta.SetDistTensorGlobalDims(static_cast<phi::distributed::DistTensor*>(
+                                         fwd_in_tensor.impl().get())
+                                         ->dims());
+        SetIsRunAutoParallel(true);
+        auto dense_tensor = static_cast<phi::distributed::DistTensor*>(
+                                fwd_in_tensor.impl().get())
+                                ->value();
+        PADDLE_ENFORCE_NE(dense_tensor.dtype(),
+                          phi::DataType::UNDEFINED,
+                          paddle::platform::errors::Fatal(
+                              "Attempting to copy DenseTensorMeta "
+                              "with phi::DataType::UNDEFINED,"
+                              "which is illegal."));
+        meta.SetTensorMeta(dense_tensor.meta());
+        meta.SetPlace(fwd_in_tensor.place());
       }
     } else {
       VLOG(7)
 
@@ -167,8 +167,20 @@ class GradSlotMeta {
 
   void SetDistAttr(const phi::distributed::TensorDistAttr& dist_attr) {
     dist_attr_ = dist_attr;
+    is_dist_meta_ = true;
   }
 
+  const phi::DDim& DistTensorGlobalDims() const {
+    return dist_tensor_global_dims_;
+  }
+
+  void SetDistTensorGlobalDims(const phi::DDim& dims) {
+    dist_tensor_global_dims_ = dims;
+    is_dist_meta_ = true;
+  }
+
+  bool IsDistMeta() const { return is_dist_meta_; }
+
  private:
   bool stop_gradient_{false};
   phi::Place place_;
@@ -178,6 +190,8 @@ class GradSlotMeta {
   // 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_;
+  phi::DDim dist_tensor_global_dims_;
+  bool is_dist_meta_{false};
 };
 
 class GradNodeBase {
 
@@ -532,12 +532,28 @@ void EagerUtils::FillZeroForEmptyOptionalGradInput(
   for (size_t i = 0; i < in_grads->size(); i++) {
     paddle::Tensor& grad = (*in_grads)[i];
     if (!grad.initialized() && grad_in_metas[i].HasTensorMeta()) {
-      auto tensor_with_zero = paddle::experimental::full(
-          phi::vectorize(grad_in_metas[i].GetTensorMeta().dims),
-          0.0,
-          grad_in_metas[i].GetTensorMeta().dtype,
-          grad_in_metas[i].GetPlace());
-      grad.set_impl(tensor_with_zero.impl());
+      if (grad_in_metas[i].IsDistMeta()) {
+        grad.set_impl(std::make_shared<phi::distributed::DistTensor>(
+            grad_in_metas[i].DistTensorGlobalDims(),
+            grad_in_metas[i].DistAttr()));
+        if (grad_in_metas[i].GetTensorMeta().dims.size() != -1) {
+          auto tensor_with_zero = paddle::experimental::full(
+              phi::vectorize(grad_in_metas[i].GetTensorMeta().dims),
+              0.0,
+              grad_in_metas[i].GetTensorMeta().dtype,
+              grad_in_metas[i].GetPlace());
+          *(static_cast<phi::distributed::DistTensor*>(grad.impl().get())
+                ->unsafe_mutable_value()) =
+              *(static_cast<phi::DenseTensor*>(tensor_with_zero.impl().get()));
+        }
+      } else {
+        auto tensor_with_zero = paddle::experimental::full(
+            phi::vectorize(grad_in_metas[i].GetTensorMeta().dims),
+            0.0,
+            grad_in_metas[i].GetTensorMeta().dtype,
+            grad_in_metas[i].GetPlace());
+        grad.set_impl(tensor_with_zero.impl());
+      }
     }
   }
 }
@@ -551,13 +567,29 @@ void EagerUtils::FillZeroForEmptyOptionalGradOutput(
       if (grad.defined() && grad.is_selected_rows()) {
         continue;
       }
-      auto tensor_with_zero =
-          paddle::experimental::full(  // only create dense tensor.
+      if (grad_output_metas[i].IsDistMeta()) {
+        grad.set_impl(std::make_shared<phi::distributed::DistTensor>(
+            grad_output_metas[i].DistTensorGlobalDims(),
+            grad_output_metas[i].DistAttr()));
+        if (grad_output_metas[i].GetTensorMeta().dims.size() != -1) {
+          auto tensor_with_zero = paddle::experimental::full(
               phi::vectorize(grad_output_metas[i].GetTensorMeta().dims),
               0.0,
               grad_output_metas[i].GetTensorMeta().dtype,
               grad_output_metas[i].GetPlace());
-      grad.set_impl(tensor_with_zero.impl());
+          *(static_cast<phi::distributed::DistTensor*>(grad.impl().get())
+                ->unsafe_mutable_value()) =
+              *(static_cast<phi::DenseTensor*>(tensor_with_zero.impl().get()));
+        }
+      } else {
+        auto tensor_with_zero =
+            paddle::experimental::full(  // only create dense tensor.
+                phi::vectorize(grad_output_metas[i].GetTensorMeta().dims),
+                0.0,
+                grad_output_metas[i].GetTensorMeta().dtype,
+                grad_output_metas[i].GetPlace());
+        grad.set_impl(tensor_with_zero.impl());
+      }
     }
   }
 }
@@ -570,25 +602,55 @@ void EagerUtils::FillZeroForEmptyGradInput(paddle::Tensor* in_grad,
         paddle::platform::errors::Fatal(
             "Unable to fill empty grad inputs due to empty GradSlotMeta"));
     const auto& tensor_meta = grad_in_meta.GetTensorMeta();
-    auto tensor_with_zero =
-        paddle::experimental::full(phi::vectorize(tensor_meta.dims),
-                                   0.0,
-                                   tensor_meta.dtype,
-                                   grad_in_meta.GetPlace());
-    in_grad->set_impl(tensor_with_zero.impl());
+    if (grad_in_meta.IsDistMeta()) {
+      in_grad->set_impl(std::make_shared<phi::distributed::DistTensor>(
+          grad_in_meta.DistTensorGlobalDims(), grad_in_meta.DistAttr()));
+      if (tensor_meta.dims.size() != -1) {
+        auto tensor_with_zero =
+            paddle::experimental::full(phi::vectorize(tensor_meta.dims),
+                                       0.0,
+                                       tensor_meta.dtype,
+                                       grad_in_meta.GetPlace());
+        *(static_cast<phi::distributed::DistTensor*>(in_grad->impl().get())
+              ->unsafe_mutable_value()) =
+            *(static_cast<phi::DenseTensor*>(tensor_with_zero.impl().get()));
+      }
+    } else {
+      auto tensor_with_zero =
+          paddle::experimental::full(phi::vectorize(tensor_meta.dims),
+                                     0.0,
+                                     tensor_meta.dtype,
+                                     grad_in_meta.GetPlace());
+      in_grad->set_impl(tensor_with_zero.impl());
+    }
   }
 }
 
 void EagerUtils::FillZeroForEmptyOptionalGradInput(
     paddle::Tensor* in_grad, const GradSlotMeta& grad_in_meta) {
   if (!in_grad->initialized() && grad_in_meta.HasTensorMeta()) {
     const auto& tensor_meta = grad_in_meta.GetTensorMeta();
-    auto tensor_with_zero =
-        paddle::experimental::full(phi::vectorize(tensor_meta.dims),
-                                   0.0,
-                                   tensor_meta.dtype,
-                                   grad_in_meta.GetPlace());
-    in_grad->set_impl(tensor_with_zero.impl());
+    if (grad_in_meta.IsDistMeta()) {
+      in_grad->set_impl(std::make_shared<phi::distributed::DistTensor>(
+          grad_in_meta.DistTensorGlobalDims(), grad_in_meta.DistAttr()));
+      if (tensor_meta.dims.size() != -1) {
+        auto tensor_with_zero =
+            paddle::experimental::full(phi::vectorize(tensor_meta.dims),
+                                       0.0,
+                                       tensor_meta.dtype,
+                                       grad_in_meta.GetPlace());
+        *(static_cast<phi::distributed::DistTensor*>(in_grad->impl().get())
+              ->unsafe_mutable_value()) =
+            *(static_cast<phi::DenseTensor*>(tensor_with_zero.impl().get()));
+      }
+    } else {
+      auto tensor_with_zero =
+          paddle::experimental::full(phi::vectorize(tensor_meta.dims),
+                                     0.0,
+                                     tensor_meta.dtype,
+                                     grad_in_meta.GetPlace());
+      in_grad->set_impl(tensor_with_zero.impl());
+    }
   }
 }