[Auto Parallel] fix bwd auto parallel (#69203)

zhangyuqin1998 · web-flow · commit 4ab178c567db · 2024-11-18T11:49:28.000+08:00
* [Auto Parallel] fix bwd auto parallel

* fix

* Update utils.h

* Update auto_parallel_backward.py

* Update multiply_node.cc
diff --git a/paddle/fluid/eager/api/manual/eager_manual/nodes/conv2d_nodes.cc b/paddle/fluid/eager/api/manual/eager_manual/nodes/conv2d_nodes.cc
@@ -28,6 +28,9 @@
 #include "paddle/fluid/eager/api/manual/eager_manual/nodes/nodes.h"
 #include "paddle/phi/api/include/sparse_api.h"
 
+using egr::ConvertAllInputsToDistTensor;
+using egr::InputsContainDistTensor;
+
 COMMON_DECLARE_bool(check_nan_inf);
 
 paddle::small_vector<std::vector<paddle::Tensor>, egr::kSlotSmallVectorSize>
@@ -62,6 +65,14 @@ Conv2dGradNodeFinal::operator()(
   auto& groups = this->groups_;
   auto& dilations = this->dilations_;
   auto& data_format = this->data_format_;
+
+  // 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, input, filter);
+  }
+
   // Prepare Grad function call
 
   const auto& out_metas = OutputMeta();
@@ -82,9 +93,9 @@ Conv2dGradNodeFinal::operator()(
           : &returns[1][0];
 
   // Set DistAttr of Out Tensor for semi-auto parallel
-  if (IsRunAutoParallel()) {
+  if (IsRunAutoParallel() || inputs_contain_dist_tensor) {
     egr::EagerUtils::SetGradOutputDistAttr(
-        out_metas, {0, 1}, api_output_0, api_output_1);
+        out_metas, {0, 1}, *mesh, api_output_0, api_output_1);
   }
 
   // Runtime check if we need next grad
diff --git a/paddle/fluid/eager/api/manual/eager_manual/nodes/multiply_node.cc b/paddle/fluid/eager/api/manual/eager_manual/nodes/multiply_node.cc
@@ -97,12 +97,9 @@ MultiplyGradNode::operator()(
   bool trace_backward = egr::Controller::Instance().HasGrad() && create_graph;
 
   // Set DistAttr of Out Tensor for semi-auto parallel
-  if (IsRunAutoParallel()) {
+  if (IsRunAutoParallel() || inputs_contain_dist_tensor) {
     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);
+        out_metas, {0, 1}, *mesh, api_output_0, api_output_1);
   }
 
   // Inplace Check
diff --git a/paddle/fluid/eager/auto_code_generator/generator/eager_gen.py b/paddle/fluid/eager/auto_code_generator/generator/eager_gen.py
@@ -285,6 +285,8 @@ class {} : public egr::GradNodeBase {{
   auto hooked_grads = ApplyGradientHooks(grads);
 
   // Collect GradIn Tensors, Attrs and Recovered TensorWrappers
+{}
+  // Convert All Inputs to DistTensor if Necessary
 {}
   // Prepare Grad function call
 {}
@@ -651,11 +653,25 @@ class {} : public egr::GradNodeBase {{
 """
 
 SET_GRAD_OUT_DIST_ATTR_TEMPLATE = """
-  if (IsRunAutoParallel()) {{
-    egr::EagerUtils::SetGradOutputDistAttr(out_metas, {}, {});
+  if (IsRunAutoParallel() || inputs_contain_dist_tensor) {{
+    egr::EagerUtils::SetGradOutputDistAttr(out_metas, {}, *mesh, {});
+  }}
+"""
+
+CONVERT_INPUT_TENSORS_TO_DIST_TENSOR_TEMPLATE = """
+  const phi::distributed::ProcessMesh* mesh = nullptr;
+  bool inputs_contain_dist_tensor = egr::InputsContainDistTensor(&mesh{grad_inputs_names});
+  if (inputs_contain_dist_tensor) {{
+    egr::ConvertAllInputsToDistTensor(mesh{wrapper_inputs_names});
   }}
 """
 
+INPUT_CONTAIN_DIST_TENSOR_TEMPLATE = """
+  const phi::distributed::ProcessMesh* mesh = nullptr;
+  bool inputs_contain_dist_tensor = false;
+  egr::InputsContainDistTensor(&mesh{grad_inputs_names});
+"""
+
 CHECK_BACKWARD_INPLACE_TEMPLATE = """
   bool can_be_inplaced = false;
   if ({}.initialized()) {{
@@ -2409,6 +2425,8 @@ def GenerateNodeDefinition(
         get_grad_in_args_list = []
         grad_api_out_args_list = []
         fwd_positions_list = []
+        grad_inputs_names = []
+        wrapper_inputs_names = []
 
         # Fill Grad Ins with Zero
         fill_zero_str = ""
@@ -2450,6 +2468,7 @@ def GenerateNodeDefinition(
         ) in backward_forward_inputs_map.items():
             tensor_wrapper_name = GetSavedName(name)
             transformed_tensor_name = self.TransformToNextGradName(name)
+            wrapper_inputs_names.append(transformed_tensor_name)
 
             is_optional = name in self.optional_inputs
             tensor_wrapper_recover_str = f"{indent}auto {transformed_tensor_name} = egr::EagerUtils::RecoverTensorWrapper(&this->{tensor_wrapper_name});"
@@ -2523,6 +2542,7 @@ def GenerateNodeDefinition(
             grad_api_position,
         ) in backward_grad_inputs_map.items():
             transformed_tensor_name = self.TransformToNextGradName(name)
+            grad_inputs_names.append(transformed_tensor_name)
 
             is_optional = name in self.optional_inputs
             if IsPlainTensorType(ttype):
@@ -2905,12 +2925,29 @@ def _gen_api_call_code_block(
 
         log_str = AFTER_LOG_PRINT_TEMPLATE.format(var_str)
 
+        if len(wrapper_inputs_names) > 0:
+            convert_input_to_dist_tensor_str = (
+                CONVERT_INPUT_TENSORS_TO_DIST_TENSOR_TEMPLATE.format(
+                    grad_inputs_names=", " + ", ".join(grad_inputs_names),
+                    wrapper_inputs_names=", " + ", ".join(wrapper_inputs_names),
+                )
+            )
+        elif not is_invoke_forward_api:
+            convert_input_to_dist_tensor_str = (
+                INPUT_CONTAIN_DIST_TENSOR_TEMPLATE.format(
+                    grad_inputs_names=", " + ", ".join(grad_inputs_names),
+                )
+            )
+        else:
+            convert_input_to_dist_tensor_str = ""
+
         self.node_definition_str = GRAD_FUNCTION_TEMPLATE.format(
             grad_node_name,
             self.backward_api_name,
             grad_node_name,
             fill_zero_str,
             get_grad_in_args_str,
+            convert_input_to_dist_tensor_str,
             grad_function_prepare_str,
             compute_require_next_grad_str,
             set_out_dist_attr_str,
diff --git a/paddle/fluid/eager/utils.cc b/paddle/fluid/eager/utils.cc
@@ -38,11 +38,25 @@ void SetGradOutputDistAttrIter::visit_element(paddle::Tensor* element,
                "SetGradOutputDistAttrIter.";
     return;
   }
-  // Here the element is empty or defined DistTensor
-  VLOG(4) << "The input element is set DistTensor impl when calling "
-             "SetGradOutputDistAttrIter.";
-  element->set_impl(std::make_shared<phi::distributed::DistTensor>(
-      phi::DDim(), meta.DistAttr()));
+  if (meta.IsDistMeta()) {
+    // Here the element is empty or defined DistTensor
+    VLOG(4) << "The input element is set DistTensor impl when calling "
+               "SetGradOutputDistAttrIter.";
+    element->set_impl(std::make_shared<phi::distributed::DistTensor>(
+        phi::DDim(), meta.DistAttr()));
+  } else {
+    // Here the element is empty or defined DenseTensor
+    VLOG(4) << "The input element is set DistTensor impl using dense meta "
+               "when calling SetGradOutputDistAttrIter.";
+    phi::distributed::Placements placements;
+    for (int64_t i = 0; i < mesh_.ndim(); ++i) {
+      placements.emplace_back(std::make_shared<phi::distributed::Replicate>());
+    }
+    auto dist_attr = phi::distributed::ToTensorDistAttr(
+        mesh_, placements, meta.GetTensorMeta().dims);
+    element->set_impl(
+        std::make_shared<phi::distributed::DistTensor>(phi::DDim(), dist_attr));
+  }
 }
 
 void SetGradOutputDistAttrIter::visit(paddle::Tensor* element) {
@@ -945,22 +959,6 @@ std::string EagerUtils::TensorStr(
   }
 }
 
-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())
diff --git a/paddle/fluid/eager/utils.h b/paddle/fluid/eager/utils.h
@@ -94,8 +94,9 @@ class SetGradOutputDistAttrIter : public IterHelper<paddle::Tensor*> {
   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} {}
+      const paddle::small_vector<size_t, kSlotSmallVectorSize>& out_indexes,
+      const phi::distributed::ProcessMesh& mesh)
+      : out_meta_(out_meta), out_indexes_{out_indexes}, mesh_(mesh) {}
 
  private:
   void visit_element(paddle::Tensor* element, const GradSlotMeta& meta);
@@ -105,6 +106,7 @@ class SetGradOutputDistAttrIter : public IterHelper<paddle::Tensor*> {
   const paddle::small_vector<std::vector<GradSlotMeta>, kSlotSmallVectorSize>&
       out_meta_;
   const paddle::small_vector<size_t, kSlotSmallVectorSize>& out_indexes_;
+  const phi::distributed::ProcessMesh& mesh_;
 
   int cur_pos_{0};
 };
@@ -264,17 +266,12 @@ class TEST_API EagerUtils {
       const paddle::small_vector<std::vector<GradSlotMeta>,
                                  kSlotSmallVectorSize>& out_metas,
       const paddle::small_vector<size_t, kSlotSmallVectorSize>& out_indexes,
+      const phi::distributed::ProcessMesh& mesh,
       Args&&... args) {
-    SetGradOutputDistAttrIter(out_metas, out_indexes)
+    SetGradOutputDistAttrIter(out_metas, out_indexes, mesh)
         .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)
    * **/
diff --git a/test/auto_parallel/CMakeLists.txt b/test/auto_parallel/CMakeLists.txt
@@ -133,9 +133,9 @@ if(WITH_DISTRIBUTE AND WITH_GPU)
     test_semi_auto_parallel_unshard_dtensor ENVS FLAGS_enable_pir_api=1)
   set_tests_properties(test_semi_auto_parallel_unshard_dtensor
                        PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 100)
-  py_test_modules(test_auto_parallel_multiply_test MODULES
-                  test_auto_parallel_multiply_test ENVS)
-  set_tests_properties(test_auto_parallel_multiply_test
+  py_test_modules(test_auto_parallel_backward_test MODULES
+                  test_auto_parallel_backward_test ENVS)
+  set_tests_properties(test_auto_parallel_backward_test
                        PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 100)
   py_test_modules(test_moe_utils MODULES test_moe_utils)
   set_tests_properties(test_moe_utils PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE"
diff --git a/test/auto_parallel/auto_parallel_backward.py b/test/auto_parallel/auto_parallel_backward.py
@@ -16,8 +16,8 @@
 import paddle.distributed as dist
 
 
-class TestMultiplyGradAutoParallel:
-    def __init__(self):
+class TestBackwardAutoParallel:
+    def init_data(self):
         self.mesh = dist.ProcessMesh([0], dim_names=['d0'])
 
         self.x = paddle.to_tensor([[1]])
@@ -30,14 +30,22 @@ def __init__(self):
 
         self.z = dist.shard_tensor(self.z, self.mesh, [dist.Replicate()])
 
-    def run_test_case(self):
+    def run_test_case1(self):
+        self.init_data()
         o = self.x * self.y
         o = o + self.z
         o = o.sum()
+        o.backward()
 
+    def run_test_case2(self):
+        self.init_data()
+        o = self.x + self.y
+        o = o - self.z
+        o = o.sum()
         o.backward()
 
 
-# python -m paddle.distributed.launch --device=0 auto_parallel_multiply_grad.py
+# python -m paddle.distributed.launch --device=0 auto_parallel_backward.py
 if __name__ == '__main__':
-    TestMultiplyGradAutoParallel().run_test_case()
+    TestBackwardAutoParallel().run_test_case1()
+    TestBackwardAutoParallel().run_test_case2()
diff --git a/test/auto_parallel/test_auto_parallel_backward_test.py b/test/auto_parallel/test_auto_parallel_backward_test.py
@@ -16,7 +16,7 @@
 import collective.test_communication_api_base as test_base
 
 
-class TestMultiplyGradAutoParallel(test_base.CommunicationTestDistBase):
+class TestBackwardAutoParallel(test_base.CommunicationTestDistBase):
     def setUp(self):
         super().setUp(
             num_of_devices=1,
@@ -31,7 +31,7 @@ def test_multiply_grad(self):
         )
         for envs in envs_list:
             self.run_test_case(
-                "auto_parallel_multiply_grad.py",
+                "auto_parallel_backward.py",
                 user_defined_envs=envs,
             )
 

Original file line number	Diff line number	Diff line change
`@@ -16,7 +16,7 @@`
`16`	`16`	`import collective.test_communication_api_base as test_base`
`17`	`17`
`18`	`18`
`19`		`-class TestMultiplyGradAutoParallel(test_base.CommunicationTestDistBase):`
	`19`	`+class TestBackwardAutoParallel(test_base.CommunicationTestDistBase):`
`20`	`20`	`def setUp(self):`
`21`	`21`	`super().setUp(`
`22`	`22`	`num_of_devices=1,`
`@@ -31,7 +31,7 @@ def test_multiply_grad(self):`
`31`	`31`	`)`
`32`	`32`	`for envs in envs_list:`
`33`	`33`	`self.run_test_case(`
`34`		`- "auto_parallel_multiply_grad.py",`
	`34`	`+ "auto_parallel_backward.py",`
`35`	`35`	`user_defined_envs=envs,`
`36`	`36`	`)`
`37`	`37`