[Auto Parallel]: Support std::vector<phi::Tensor> input and output for DistTensor. (#56602)

GhostScreaming · web-flow · commit d2fedeac0a35 · 2023-09-05T14:32:57.000+08:00
* [WIP] Support std::vector&lt;phi::Tensor&gt; input and output for DistTensor.
Concat forward and backward are verified.

* Polish code for new dist tensor implementation.

* Fix bug of DistTensor upgrade. Add support functions for std::vector&lt;Tensor&gt; -&gt; std::vector&lt;Tensor&gt;.

* Add support for DistTensor type of std::vector&lt;phi::Tensor&gt; as input or output of operators.
Following testcases are passed.
1. concat: std::vector&lt;phi::Tensor&gt; -&gt; phi::Tensor
2. unbind: phi::Tensor -&gt; std::vector&lt;phi::Tensor&gt;
3. broadcast_tensors: std::vector&lt;phi::Tensor&gt; -&gt; std::vector&lt;phi::Tensor&gt;

* Polish code. Remove useless comments.

* Add update_loss_scaling in skip_op_lists.

* Polish code.
diff --git a/paddle/fluid/eager/tensor_wrapper.h b/paddle/fluid/eager/tensor_wrapper.h
@@ -32,6 +32,8 @@
 #ifndef PADDLE_NO_PYTHON
 #include "paddle/fluid/eager/hooks.h"
 #endif
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
 
 namespace egr {
 class TensorWrapper {
@@ -66,6 +68,16 @@ class TensorWrapper {
         intermidiate_tensor_.set_impl(std::make_shared<phi::DenseTensor>(
             std::make_shared<phi::Allocation>(nullptr, 0, tensor.place()),
             dense_tensor->meta()));
+      } else if (phi::distributed::DistTensor::classof(tensor.impl().get())) {
+        // Only Copy Meta
+        phi::distributed::DistTensor* dist_tensor =
+            static_cast<phi::distributed::DistTensor*>(tensor.impl().get());
+        intermidiate_tensor_.set_impl(
+            std::make_shared<phi::distributed::DistTensor>(
+                phi::DenseTensor(std::make_shared<phi::Allocation>(
+                                     nullptr, 0, tensor.place()),
+                                 dist_tensor->value().meta()),
+                dist_tensor->dist_attr()));
       } else {
         PADDLE_THROW(paddle::platform::errors::Fatal(
             "Unrecognized tensor type for no_need_buffer feature"));
diff --git a/paddle/phi/api/lib/api_gen_utils.cc b/paddle/phi/api/lib/api_gen_utils.cc
@@ -192,8 +192,6 @@ std::vector<phi::MetaTensor> MakeMetaTensor(
   return meta_tensors;
 }
 
-/* ------------------ for output ----------------------- */
-
 phi::DenseTensor* SetKernelOutput(Tensor* out) {
   if (out) {
     if (out->impl() == nullptr) {
@@ -546,5 +544,45 @@ phi::distributed::DistTensor* SetKernelDistOutput(Tensor* out) {
   return nullptr;
 }
 
+std::vector<phi::distributed::DistTensor*> SetKernelDistOutput(
+    std::vector<Tensor*> out) {
+  std::vector<phi::distributed::DistTensor*> result;
+  for (auto tmp : out) {
+    if (tmp) {
+      // TODO(GhostScreaming): now all dist case are nullptr
+      if (tmp->impl() == nullptr) {
+        phi::DenseTensor dense_t;
+        // TODO(GhostScreaming): polish code, dist_attr is null now
+        phi::distributed::TensorDistAttr dist_attr;
+        auto dist_t =
+            std::make_shared<phi::distributed::DistTensor>(dense_t, dist_attr);
+        tmp->set_impl(dist_t);
+      }
+      result.emplace_back(
+          static_cast<phi::distributed::DistTensor*>(tmp->impl().get()));
+    } else {
+      result.emplace_back(nullptr);
+    }
+  }
+  return result;
+}
+
+std::vector<phi::distributed::DistTensor*> SetKernelDistOutput(
+    size_t out_size, std::vector<Tensor>* out) {
+  out->reserve(out_size);
+  std::vector<phi::distributed::DistTensor*> results(out_size);
+  for (size_t i = 0; i < out_size; ++i) {
+    phi::DenseTensor dense_t;
+    // TODO(GhostScreaming): polish code, dist_attr is null now
+    phi::distributed::TensorDistAttr dist_attr;
+    auto dist_t =
+        std::make_shared<phi::distributed::DistTensor>(dense_t, dist_attr);
+    results[i] = dist_t.get();
+    out->emplace_back();
+    out->back().set_impl(dist_t);
+  }
+  return results;
+}
+
 }  // namespace experimental
 }  // namespace paddle
diff --git a/paddle/phi/api/lib/api_gen_utils.h b/paddle/phi/api/lib/api_gen_utils.h
@@ -140,6 +140,10 @@ void TransStrideLegacy(phi::DeviceContext* dev_ctx,
 /* ------------------ for auto parallel ----------------------- */
 
 phi::distributed::DistTensor* SetKernelDistOutput(Tensor* out);
+std::vector<phi::distributed::DistTensor*> SetKernelDistOutput(
+    std::vector<Tensor*> out);
+std::vector<phi::distributed::DistTensor*> SetKernelDistOutput(
+    size_t out_size, std::vector<Tensor>* out);
 
 }  // namespace experimental
 }  // namespace paddle
diff --git a/paddle/phi/api/lib/data_transform.cc b/paddle/phi/api/lib/data_transform.cc
@@ -632,5 +632,47 @@ std::shared_ptr<phi::distributed::DistTensor> PrepareDataForDistTensor(
   return nullptr;
 }
 
+std::vector<std::shared_ptr<phi::distributed::DistTensor>>
+PrepareDataForDistTensor(const std::vector<Tensor>& input,
+                         const phi::TensorArgDef& target_args_def,
+                         const TransformFlag& transform_flag,
+                         bool is_stride_kernel) {
+  std::vector<std::shared_ptr<phi::distributed::DistTensor>> out;
+  for (auto x : input) {
+    const auto& tensor_in = x.impl();
+    if (tensor_in) {
+      phi::distributed::DistTensor* dist_tensor =
+          static_cast<phi::distributed::DistTensor*>(tensor_in.get());
+      const phi::DenseTensor& dense_tensor = dist_tensor->value();
+      if (!transform_flag.NeedTransform() || !dense_tensor.initialized() ||
+          (!NeedTransformPlace(
+               dense_tensor.place(), target_args_def.backend, transform_flag) &&
+           !NeedTransformDataType(
+               dense_tensor.dtype(), target_args_def.dtype, transform_flag) &&
+           !NeedTransformLayout(dense_tensor.layout(),
+                                target_args_def.layout,
+                                dense_tensor.place(),
+                                transform_flag) &&
+           !NeedTransform2Contiguous(is_stride_kernel,
+                                     dense_tensor.meta().is_contiguous()))) {
+        out.push_back(
+            std::static_pointer_cast<phi::distributed::DistTensor>(tensor_in));
+        continue;
+      }
+      phi::DenseTensor trans_in_tensor = TransformData(
+          dense_tensor, target_args_def, transform_flag, is_stride_kernel);
+      // TODO(GhostScreaming): The global meta in DistTensor is not changed,
+      // but the local meta in DenseTensor maybe changed, such as layout
+      // change(NCHW->NHWC), so the new DistTensor's meta maybe not unified.
+      VLOG(6) << "PrepareDataForDistTensor return transformed dist tensor";
+      out.push_back(std::make_shared<phi::distributed::DistTensor>(
+          trans_in_tensor, dist_tensor->dist_attr()));
+    } else {
+      out.push_back(nullptr);
+    }
+  }
+  return out;
+}
+
 }  // namespace experimental
 }  // namespace paddle
diff --git a/paddle/phi/api/lib/data_transform.h b/paddle/phi/api/lib/data_transform.h
@@ -180,5 +180,11 @@ std::shared_ptr<phi::distributed::DistTensor> PrepareDataForDistTensor(
     const TransformFlag& transform_flag,
     bool is_stride_kernel);
 
+std::vector<std::shared_ptr<phi::distributed::DistTensor>>
+PrepareDataForDistTensor(const std::vector<Tensor>& input,
+                         const phi::TensorArgDef& target_args_def,
+                         const TransformFlag& transform_flag,
+                         bool is_stride_kernel);
+
 }  // namespace experimental
 }  // namespace paddle
diff --git a/paddle/phi/api/lib/kernel_dispatch.h b/paddle/phi/api/lib/kernel_dispatch.h
@@ -191,6 +191,16 @@ struct DistTensorTypeParser : ArgsIterator<DistTensorTypeParser> {
     }
   }
 
+  void operator()(const paddle::optional<std::vector<Tensor>>& x) {
+    if (x) {
+      if (!(x.get_ptr()->empty())) {
+        for (auto& t : *(x.get_ptr())) {
+          result &= t.is_dist_tensor();
+        }
+      }
+    }
+  }
+
   // skip other type args, these args don't used in kernel selection
   template <typename T>
   void operator()(const T& x) {
diff --git a/paddle/phi/api/yaml/generator/dist_api_gen.py b/paddle/phi/api/yaml/generator/dist_api_gen.py
diff --git a/paddle/phi/api/yaml/generator/dist_bw_api_gen.py b/paddle/phi/api/yaml/generator/dist_bw_api_gen.py
diff --git a/test/auto_parallel/test_dist_tensor.py b/test/auto_parallel/test_dist_tensor.py
