code-gen backward support vector of dtensor grad output (#59328)

liuzhenhai93 · web-flow · commit ed516328f139 · 2023-11-28T16:57:43.000+08:00
* polish

* polish
diff --git a/paddle/phi/api/lib/api_gen_utils.cc b/paddle/phi/api/lib/api_gen_utils.cc
@@ -641,14 +641,10 @@ std::vector<phi::distributed::DistTensor*> SetKernelDistOutput(
 std::shared_ptr<phi::distributed::DistTensor> CreateKernelDistOutput(
     Tensor* out,
     bool set_dist_output_as_tensor_impl,
-    const phi::distributed::ArgDistAttr& dist_attr) {
+    const phi::distributed::TensorDistAttr& dist_attr) {
   if (out) {
-    PADDLE_ENFORCE_EQ(
-        paddle::holds_alternative<phi::distributed::TensorDistAttr>(dist_attr),
-        true,
-        phi::errors::PreconditionNotMet("Arg must be a single TensorDistAttr"));
-    auto dist_output = std::make_shared<phi::distributed::DistTensor>(
-        phi::DDim(), paddle::get<0>(dist_attr));
+    auto dist_output =
+        std::make_shared<phi::distributed::DistTensor>(phi::DDim(), dist_attr);
     if (set_dist_output_as_tensor_impl) {
       VLOG(3) << "CreateKernelDistOutput function set generated output "
                  "dist_tensor as Tensor's impl";
@@ -666,17 +662,57 @@ std::shared_ptr<phi::distributed::DistTensor> CreateKernelDistOutput(
   return nullptr;
 }
 
+std::shared_ptr<phi::distributed::DistTensor> CreateKernelDistOutput(
+    Tensor* out,
+    bool set_dist_output_as_tensor_impl,
+    const phi::distributed::ArgDistAttr& dist_attr) {
+  auto& tensor_dist_attr =
+      PADDLE_GET_CONST(phi::distributed::TensorDistAttr, dist_attr);
+  return CreateKernelDistOutput(
+      out, set_dist_output_as_tensor_impl, tensor_dist_attr);
+}
+
 std::shared_ptr<phi::distributed::DistTensor> CreateKernelDistOutput(
     Tensor* out, const phi::distributed::ArgDistAttr& dist_attr) {
-  if (out) {
-    PADDLE_ENFORCE_EQ(
-        paddle::holds_alternative<phi::distributed::TensorDistAttr>(dist_attr),
-        true,
-        phi::errors::PreconditionNotMet("Arg must be a single TensorDistAttr"));
-    return std::make_shared<phi::distributed::DistTensor>(
-        phi::DDim(), paddle::get<0>(dist_attr));
+  auto& tensor_dist_attr =
+      PADDLE_GET_CONST(phi::distributed::TensorDistAttr, dist_attr);
+  return CreateKernelDistOutput(out, false, tensor_dist_attr);
+}
+
+std::vector<std::shared_ptr<phi::distributed::DistTensor>>
+CreateKernelDistOutput(std::vector<Tensor*> out,
+                       bool set_dist_output_as_tensor_impl,
+                       const phi::distributed::ArgDistAttr& dist_attr) {
+  auto tensor_dist_attrs = PADDLE_GET_CONST(
+      std::vector<phi::distributed::TensorDistAttr>, dist_attr);
+  PADDLE_ENFORCE_EQ(
+      out.size(),
+      tensor_dist_attrs.size(),
+      phi::errors::PreconditionNotMet(
+          "out.size() [%d] and tensor_dist_attrs.size() [%d] not match",
+          out.size(),
+          tensor_dist_attrs.size()));
+  auto size = tensor_dist_attrs.size();
+  std::vector<std::shared_ptr<phi::distributed::DistTensor>> results;
+  results.reserve(size);
+  for (size_t i = 0; i < size; i++) {
+    results.emplace_back(CreateKernelDistOutput(
+        out[i], set_dist_output_as_tensor_impl, tensor_dist_attrs[i]));
   }
-  return nullptr;
+  return results;
+}
+
+std::vector<std::shared_ptr<phi::distributed::DistTensor>>
+CreateKernelDistOutput(std::vector<Tensor*> out,
+                       bool set_dist_output_as_tensor_impl) {
+  auto size = out.size();
+  std::vector<std::shared_ptr<phi::distributed::DistTensor>> results;
+  results.reserve(size);
+  for (size_t i = 0; i < size; i++) {
+    results.emplace_back(
+        CreateKernelDistOutput(out[i], set_dist_output_as_tensor_impl));
+  }
+  return results;
 }
 
 void SetReplicatedDistAttrForOutput(
diff --git a/paddle/phi/api/lib/api_gen_utils.h b/paddle/phi/api/lib/api_gen_utils.h
@@ -158,12 +158,26 @@ std::vector<phi::distributed::DistTensor*> SetKernelDistOutput(
 std::vector<phi::distributed::DistTensor*> SetKernelDistOutput(
     const phi::distributed::ArgDistAttr& dist_attr, std::vector<Tensor>* out);
 
+std::shared_ptr<phi::distributed::DistTensor> CreateKernelDistOutput(
+    Tensor* out,
+    bool set_dist_output_as_tensor_impl,
+    const phi::distributed::TensorDistAttr& dist_attr);
+
 std::shared_ptr<phi::distributed::DistTensor> CreateKernelDistOutput(
     Tensor* out,
     bool set_dist_output_as_tensor_impl,
     const phi::distributed::ArgDistAttr& dist_attr =
         phi::distributed::TensorDistAttr());
 
+std::vector<std::shared_ptr<phi::distributed::DistTensor>>
+CreateKernelDistOutput(std::vector<Tensor*> out,
+                       bool set_dist_output_as_tensor_impl,
+                       const phi::distributed::ArgDistAttr& dist_attr);
+
+std::vector<std::shared_ptr<phi::distributed::DistTensor>>
+CreateKernelDistOutput(std::vector<Tensor*> out,
+                       bool set_dist_output_as_tensor_impl);
+
 std::shared_ptr<phi::distributed::DistTensor> CreateKernelDistOutput(
     Tensor* out, const phi::distributed::ArgDistAttr& dist_attr);
 
diff --git a/paddle/phi/api/lib/data_transform.cc b/paddle/phi/api/lib/data_transform.cc
@@ -894,6 +894,24 @@ void ReshardKernelOutputToApiOutput(
   }
 }
 
+void ReshardKernelOutputToApiOutput(
+    phi::DeviceContext* dev_ctx,
+    const std::vector<std::shared_ptr<phi::distributed::DistTensor>>&
+        src_tensors,
+    const std::vector<Tensor*>& dst_tensors) {
+  PADDLE_ENFORCE_EQ(
+      src_tensors.size(),
+      dst_tensors.size(),
+      phi::errors::PreconditionNotMet(
+          "src_tensors.size() [%d] and dst_tensors.size() [%d] not match",
+          src_tensors.size(),
+          dst_tensors.size()));
+  auto size = src_tensors.size();
+  for (size_t i = 0; i < size; i++) {
+    ReshardKernelOutputToApiOutput(dev_ctx, src_tensors[i], dst_tensors[i]);
+  }
+}
+
 std::shared_ptr<phi::distributed::DistTensor> PrepareDataForDistTensor(
     std::shared_ptr<phi::distributed::DistTensor> input,
     const phi::TensorArgDef& target_args_def,
diff --git a/paddle/phi/api/lib/data_transform.h b/paddle/phi/api/lib/data_transform.h
@@ -236,6 +236,12 @@ void ReshardKernelOutputToApiOutput(
     const std::shared_ptr<phi::distributed::DistTensor>& src_tensor,
     Tensor* dst_tensor);
 
+void ReshardKernelOutputToApiOutput(
+    phi::DeviceContext* dev_ctx,
+    const std::vector<std::shared_ptr<phi::distributed::DistTensor>>&
+        src_tensors,
+    const std::vector<Tensor*>& dst_tensors);
+
 std::shared_ptr<phi::distributed::DistTensor> PrepareDataForDistTensor(
     std::shared_ptr<phi::distributed::DistTensor> input,
     const phi::TensorArgDef& target_args_def,
diff --git a/paddle/phi/api/yaml/generator/dist_bw_api_gen.py b/paddle/phi/api/yaml/generator/dist_bw_api_gen.py
@@ -71,7 +71,7 @@
             phi::DenseTensorMeta());
     }}
 """
-VECTOR_OUT_CREATION_TEMPLATE = """
+VECTOR_OUT_CREATION_TEMPLATE_WITH_NO_SPMD = """
     auto dist_out = SetKernelDistOutput({name});
     std::vector<phi::DenseTensor*> dense_out(dist_out.size());
     for (size_t i=0; i<dist_out.size(); i++) {{
@@ -83,6 +83,43 @@
       }}
     }}
 """
+
+VECTOR_OUT_CREATION_TEMPLATE_WITH_SPMD = """
+    auto shared_dist_out = CreateKernelDistOutput({name}, !rank_is_in_current_mesh, spmd_info.second[0]);
+    std::vector<phi::distributed::DistTensor*> dist_out;
+    for(auto& e: shared_dist_out){{
+      dist_out.push_back(e.get());
+    }}
+    std::vector<phi::DenseTensor*> dense_out(dist_out.size());
+    for (size_t i=0; i<dist_out.size(); i++) {{
+      dense_out[i] = dist_out[i]->unsafe_mutable_value();
+      if (dense_out[i] && !rank_is_in_current_mesh && !dist_out[i]->defined()) {{
+        *dense_out[i] = phi::DenseTensor(
+              std::make_shared<phi::Allocation>(nullptr, 0, phi::distributed::GetDefaultPlace()),
+              phi::DenseTensorMeta());
+      }}
+    }}
+"""
+
+
+VECTOR_OUT_CREATION_TEMPLATE = """
+    auto shared_dist_out = CreateKernelDistOutput({name}, !rank_is_in_current_mesh);
+    std::vector<phi::distributed::DistTensor*> dist_out;
+    for(auto& e: shared_dist_out){{
+      dist_out.push_back(e.get());
+    }}
+    std::vector<phi::DenseTensor*> dense_out(dist_out.size());
+    for (size_t i=0; i<dist_out.size(); i++) {{
+      dense_out[i] = dist_out[i]->unsafe_mutable_value();
+      if (dense_out[i] && !rank_is_in_current_mesh && !dist_out[i]->defined()) {{
+        *dense_out[i] = phi::DenseTensor(
+              std::make_shared<phi::Allocation>(nullptr, 0, phi::distributed::GetDefaultPlace()),
+              phi::DenseTensorMeta());
+      }}
+    }}
+"""
+
+
 INPLACE_OUT_CREATION_TEMPLATE = """
     *{} = {};
 """
@@ -133,9 +170,13 @@
 # 9. Reshard Output
 RESHARD_SINGLE_OUTPUT_TEMPLATE = """
       ReshardKernelOutputToApiOutput(dev_ctx, shared_dist_out, {});"""
+
 RESHARD_MULTI_SINGLE_OUTPUT_TEMPLATE = """
       ReshardKernelOutputToApiOutput(dev_ctx, shared_dist_out_{}, {});"""
 
+RESHARD_VECTOR_OUTPUT_TEMPLATE = """
+      ReshardKernelOutputToApiOutput(dev_ctx, shared_dist_out, {});"""
+
 
 class DistBackwardAPI(DistForwardAPI, BackwardAPI):
     def __init__(self, backward_item_yaml):
@@ -169,9 +210,22 @@ def generate_output_creation_code(self) -> str:
                         )
                     )
             elif self.outputs['types'][0] == 'std::vector<Tensor>':
-                output_creation_code += VECTOR_OUT_CREATION_TEMPLATE.format(
-                    name=self.outputs['names'][0]
-                )
+                if self.infer_meta['spmd_rule'] is not None:
+                    output_creation_code += (
+                        VECTOR_OUT_CREATION_TEMPLATE_WITH_SPMD.format(
+                            name=self.outputs['names'][0]
+                        )
+                    )
+                elif self.generate_general_infer_spmd is True:
+                    output_creation_code += VECTOR_OUT_CREATION_TEMPLATE.format(
+                        name=self.outputs['names'][0]
+                    )
+                else:
+                    output_creation_code += (
+                        VECTOR_OUT_CREATION_TEMPLATE_WITH_NO_SPMD.format(
+                            name=self.outputs['names'][0]
+                        )
+                    )
             else:
                 self.vector_output_size_assertion_check()
         elif output_num > 1:
@@ -267,6 +321,12 @@ def generate_reshard_output_code(self):
                             self.outputs['names'][0]
                         )
                     )
+                elif self.outputs['types'][0] == 'std::vector<Tensor>':
+                    reshard_output_code += (
+                        RESHARD_VECTOR_OUTPUT_TEMPLATE.format(
+                            self.outputs['names'][0]
+                        )
+                    )
                 else:
                     self.vector_output_size_assertion_check()
             elif output_num > 1:
diff --git a/test/auto_parallel/semi_auto_parallel_for_concat.py b/test/auto_parallel/semi_auto_parallel_for_concat.py
@@ -38,7 +38,7 @@ def test_concat_forward(self):
             inputs_shape=shapes,
             inputs_specs=specs,
             op_func=paddle.concat,
-            with_backward=False,
+            with_backward=True,
             axis=0,
         )
         self.check_dim_mapping(outputs, [-1, -1, 0])
@@ -50,7 +50,7 @@ def test_concat_forward_reshard(self):
             inputs_shape=shapes,
             inputs_specs=specs,
             op_func=paddle.concat,
-            with_backward=False,
+            with_backward=True,
             axis=0,
         )
         self.check_dim_mapping(outputs, [-1, -1, 0])
@@ -62,7 +62,7 @@ def test_stack_forward(self):
             inputs_shape=shapes,
             inputs_specs=specs,
             op_func=paddle.stack,
-            with_backward=False,
+            with_backward=True,
             axis=0,
         )
         self.check_dim_mapping(outputs, [-1, -1, -1, 0])
@@ -74,7 +74,7 @@ def test_stack_forward_reshard(self):
             inputs_shape=shapes,
             inputs_specs=specs,
             op_func=paddle.stack,
-            with_backward=False,
+            with_backward=True,
             axis=0,
         )
         self.check_dim_mapping(outputs, [-1, 0, -1, -1])
