[AutoParallel] Remove p2r in forward api (#58644)

chenwhql · LiYuRio · web-flow · commit 3bbb89c98200 · 2023-11-13T18:17:00.000+08:00
* remove p2r in forward api

* support reduce mean in r to p

* test only

---------

Co-authored-by: LiYuRio &lt;liyuruijx@163.com&gt;
diff --git a/paddle/fluid/pybind/eager_method.cc b/paddle/fluid/pybind/eager_method.cc
@@ -109,6 +109,7 @@ phi::DenseTensor ReshardXToReplicated(
     phi::distributed::TensorDistAttr dist_attr(dist_tensor->dist_attr());
     std::vector<int64_t> dims_mapping(dist_tensor->dims().size(), -1);
     dist_attr.set_dims_mapping(dims_mapping);
+    dist_attr.clean_partial_status();
 
     // reshard to replicate dist tensor
     auto* func =
diff --git a/paddle/phi/api/yaml/generator/dist_api_gen.py b/paddle/phi/api/yaml/generator/dist_api_gen.py
@@ -55,9 +55,8 @@
       // 7. Infer Local DenseTensor Meta{}
       // 8. DenseTensor Kernel Call{}
     }}\n
-    // 9. Reshard Partial Output to Replicated (Temporary){}\n
-    // 10. Set Output Dist Attr For Default Impl{}\n
-    // 11. Return
+    // 9. Set Output Dist Attr For Default Impl{}\n
+    // 10. Return
     {}
   }}
 """
@@ -361,17 +360,7 @@
 PREFIX_VECTOR_TENSOR_NAME = "dense_input_"
 SUFFIX_VECTOR_TENSOR_NAME = "_vec"
 
-# 9. Reshard Partial Output to Replicated
-RESHARD_P2R_SINGLE_OUTPUT_TEMPLATE = """
-    dev_ctx = phi::distributed::GetDistTensorDeviceContext(dist_out);
-    ReshardOutputPartialAxisToReplicated(dev_ctx, dist_out);"""
-RESHARD_P2R_MULTI_SINGLE_OUTPUT_TEMPLATE = """
-    dev_ctx = phi::distributed::GetDistTensorDeviceContext(dist_out_{idx});
-    ReshardOutputPartialAxisToReplicated(dev_ctx, dist_out_{idx});"""
-UNSUPPORTED_RESHARD_OUTPUT_COMMENT_TEMPLATE = """
-      // API `{}` does not need to support ReshardOutput now."""
-
-# 10. Set Output DistAttr for Default impl
+# 9. Set Output DistAttr for Default impl
 # Dist Branch will not generated in the API that doesn't have input tensor.
 CURRENT_PROCESS_MESH_TEMPLATE = """
     auto current_process_mesh = paddle::holds_alternative<phi::distributed::TensorDistAttr>(spmd_info.first[0]) ?
@@ -1363,36 +1352,6 @@ def generate_kernel_call_code(self) -> str:
                     result += MULTI_SINGLE_SET_DIST_OUT_DIMS.format(i, i)
         return result
 
-    def generate_reshard_partial_out_to_replicated_code(self) -> str:
-        reshard_p2r_code = ""
-        if self.infer_meta['spmd_rule'] is not None:
-            output_num = len(self.outputs['types'])
-            if output_num == 1:
-                if self.outputs['types'][0] == 'Tensor':
-                    reshard_p2r_code += RESHARD_P2R_SINGLE_OUTPUT_TEMPLATE
-                else:
-                    self.vector_output_size_assertion_check()
-            elif output_num > 1:
-                for i, out_type in enumerate(self.outputs['types']):
-                    if out_type == 'Tensor':
-                        reshard_p2r_code += (
-                            RESHARD_P2R_MULTI_SINGLE_OUTPUT_TEMPLATE.format(
-                                idx=i
-                            )
-                        )
-                    else:
-                        self.vector_output_size_assertion_check()
-            else:
-                raise ValueError(
-                    f"{self.api} : Output error: the output should not be empty."
-                )
-        else:
-            reshard_p2r_code = (
-                UNSUPPORTED_RESHARD_OUTPUT_COMMENT_TEMPLATE.format(self.api)
-            )
-
-        return reshard_p2r_code
-
     def generate_output_dist_attr_setting(self) -> str:
         set_out_dist_attr_code = ""
         if self.generate_general_infer_spmd is True:
@@ -1432,7 +1391,6 @@ def generate_auto_paralel_branch(self) -> str:
             self.generate_prepare_data_code(),
             self.generate_infer_meta_code(),
             self.generate_kernel_call_code(),
-            self.generate_reshard_partial_out_to_replicated_code(),
             self.generate_output_dist_attr_setting(),
             self.generate_return_code(),
         )
diff --git a/paddle/phi/api/yaml/generator/dist_bw_api_gen.py b/paddle/phi/api/yaml/generator/dist_bw_api_gen.py
@@ -36,7 +36,7 @@
       // 8. Infer Local DenseTensor Meta{}
       // 9. DenseTensor Kernel Call{}
     }}
-    // 10. Reshard Partial Output to Replicated (Temporary){}\n
+    // 10. Reshard Kernel Output to API output{}\n
     // 11. Return
     {}
   }}
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/r_to_p_reshard_function.cc b/paddle/phi/core/distributed/auto_parallel/reshard/r_to_p_reshard_function.cc
@@ -49,11 +49,19 @@ void RToPReshardFunction::Eval(phi::DeviceContext* dev_ctx,
   VLOG(3) << "Call RToPReshardFunction Eval";
   const auto& out_process_mesh = out_dist_attr.process_mesh();
   int64_t local_rank = GetCurRankCoordInMesh(out_process_mesh)[0];
-  IntArray shape(in.dims().Get(), in.dims().size());
+  const auto& in_reduce_type = out_dist_attr.partial_status().at(0);
 
   if (local_rank != 0) {
-    // reset the physical tensor to zero
-    RESHARD_FUNCTOR(dev_ctx, Full, in.dtype(), shape, 0, GetMutableTensor(out));
+    if (in_reduce_type == ReduceType::kRedAvg) {
+      // assign the input value to output
+      RESHARD_FUNCTOR_WITHOUT_DTYPE(
+          dev_ctx, Assign, in.value(), GetMutableTensor(out));
+    } else {
+      // reset the physical tensor to zero
+      IntArray shape(in.local_dims().Get(), in.local_dims().size());
+      RESHARD_FUNCTOR(
+          dev_ctx, Full, in.dtype(), shape, 0, GetMutableTensor(out));
+    }
   } else {
     // assign the input value to output
     RESHARD_FUNCTOR_WITHOUT_DTYPE(
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h b/paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h
@@ -75,14 +75,14 @@ CommContext* CreateOrGetCommContext(const DeviceContext& dev_ctx,
 #define RESHARD_FUNCTOR_IMPL(dev_ctx, fn_name, dtype, ...)            \
   do {                                                                \
     if (phi::CPUContext::classof(dev_ctx)) {                          \
-      VLOG(4) << "Call `" << #fn_name << "` in Resharding on GPU.";   \
+      VLOG(4) << "Call `" << #fn_name << "` in Resharding on CPU.";   \
       PD_VISIT_BOOL_AND_FLOATING_AND_INTEGRAL_TYPES(                  \
           dtype, #fn_name, ([&] {                                     \
             fn_name<data_t>(static_cast<const CPUContext&>(*dev_ctx), \
                             __VA_ARGS__);                             \
           }));                                                        \
     } else if (phi::GPUContext::classof(dev_ctx)) {                   \
-      VLOG(4) << "Call `" << #fn_name << "` in Resharding on CPU.";   \
+      VLOG(4) << "Call `" << #fn_name << "` in Resharding on GPU.";   \
       PD_VISIT_BOOL_AND_FLOATING_AND_INTEGRAL_TYPES(                  \
           dtype, #fn_name, ([&] {                                     \
             fn_name<data_t>(static_cast<const GPUContext&>(*dev_ctx), \
diff --git a/paddle/phi/infermeta/spmd_rules/elementwise.cc b/paddle/phi/infermeta/spmd_rules/elementwise.cc
@@ -224,8 +224,8 @@ SpmdInfo ElementwiseBinaryInferSpmd(const DistMetaTensor& x,
   out_dist_attr.set_dims_mapping(out_dims_mapping);
 
   // Step2.3: Update inputs' dims mapping with merged one.
-  TensorDistAttr x_dist_attr_dst(x_dist_attr_src);
-  TensorDistAttr y_dist_attr_dst(y_dist_attr_src);
+  TensorDistAttr x_dist_attr_dst = CopyTensorDistAttrForOutput(x_dist_attr_src);
+  TensorDistAttr y_dist_attr_dst = CopyTensorDistAttrForOutput(y_dist_attr_src);
   x_dist_attr_dst.set_dims_mapping(
       GetDimsMappingForAxes(x_axes, axis_to_dim_map));
   y_dist_attr_dst.set_dims_mapping(
diff --git a/paddle/phi/infermeta/spmd_rules/matmul.cc b/paddle/phi/infermeta/spmd_rules/matmul.cc
@@ -310,6 +310,22 @@ SpmdInfo MatmulGradInferSpmd(const DistMetaTensor& x,
                                  y.dist_attr()));
   };
 
+  auto confirm_dist_attr_with_arg_same_fn = [&](const ArgDistAttr& x_dist_attr,
+                                                const ArgDistAttr& y_dist_attr,
+                                                const char* debug_msg) {
+    const auto& x_single_dist_attr = get_attr(x_dist_attr);
+    const auto& y_single_dist_attr = get_attr(y_dist_attr);
+    PADDLE_ENFORCE_EQ(
+        DistAttrsAreBasicallyEqual(x_single_dist_attr, y_single_dist_attr),
+        true,
+        phi::errors::Unavailable("The matmul grad infer spmd `%s` verify "
+                                 "error: left dist attr is %s, "
+                                 "right dist attr is %s.",
+                                 debug_msg,
+                                 x_single_dist_attr,
+                                 y_single_dist_attr));
+  };
+
   // TODO(chenweihang): Now for the case where the forward input generates
   // an intermediate value through Reshard, because the intermediate value
   // is destroyed after the forward calculation is completed, the x and y
@@ -343,6 +359,9 @@ SpmdInfo MatmulGradInferSpmd(const DistMetaTensor& x,
       confirm_dist_attr_same_fn(
           dy_spmd_info.first[0], out_grad, "trans x&y: dy-out_grad");
       confirm_dist_attr_same_fn(dy_spmd_info.first[1], x, "trans x&y: dy-x");
+      return {
+          {dy_spmd_info.first[1], dx_spmd_info.first[0], dx_spmd_info.first[1]},
+          {dx_spmd_info.second[0], dy_spmd_info.second[0]}};
     } else {
       // X'Y: dX = YG', dY = XG
       dx_spmd_info =
@@ -355,6 +374,9 @@ SpmdInfo MatmulGradInferSpmd(const DistMetaTensor& x,
       confirm_dist_attr_same_fn(dy_spmd_info.first[0], x, "trans x: dy-x");
       confirm_dist_attr_same_fn(
           dy_spmd_info.first[1], out_grad, "trans x: dy-out_grad");
+      return {
+          {dy_spmd_info.first[0], dx_spmd_info.first[0], dx_spmd_info.first[1]},
+          {dx_spmd_info.second[0], dy_spmd_info.second[0]}};
     }
   } else {
     if (trans_y) {
@@ -369,25 +391,25 @@ SpmdInfo MatmulGradInferSpmd(const DistMetaTensor& x,
       confirm_dist_attr_same_fn(
           dy_spmd_info.first[0], out_grad, "trans y: dy-out_grad");
       confirm_dist_attr_same_fn(dy_spmd_info.first[1], x, "trans y: dy-x");
+      return {
+          {dy_spmd_info.first[1], dx_spmd_info.first[1], dx_spmd_info.first[0]},
+          {dx_spmd_info.second[0], dy_spmd_info.second[0]}};
     } else {
       // XY: dX = GY', dY = X'G
       dx_spmd_info =
           MatmulInferSpmd(out_grad, y, /*trans_x=*/false, /*trans_y=*/true);
       dy_spmd_info =
           MatmulInferSpmd(x, out_grad, /*trans_x=*/true, /*trans_y=*/false);
-      confirm_dist_attr_same_fn(
-          dx_spmd_info.first[0], out_grad, "no trans: dx-out_grad");
       confirm_dist_attr_same_fn(dx_spmd_info.first[1], y, "no trans: dx-y");
       confirm_dist_attr_same_fn(dy_spmd_info.first[0], x, "no trans: dy-x");
-      confirm_dist_attr_same_fn(
-          dy_spmd_info.first[1], out_grad, "no trans: dy-out_grad");
+      confirm_dist_attr_with_arg_same_fn(dx_spmd_info.first[0],
+                                         dy_spmd_info.first[1],
+                                         "no trans: dy-out_grad");
+      return {
+          {dy_spmd_info.first[0], dx_spmd_info.first[1], dx_spmd_info.first[0]},
+          {dx_spmd_info.second[0], dy_spmd_info.second[0]}};
     }
   }
-
-  // Here we assume that the input dist attr is unchanged after inference,
-  // and only return the gradient dist attr
-  return {{x.dist_attr(), y.dist_attr(), out_grad.dist_attr()},
-          {dx_spmd_info.second[0], dy_spmd_info.second[0]}};
 }
 
 }  // namespace distributed
diff --git a/paddle/phi/infermeta/spmd_rules/reduction.cc b/paddle/phi/infermeta/spmd_rules/reduction.cc
@@ -210,8 +210,10 @@ SpmdInfo ReductionGradInferSpmd(const DistMetaTensor& x,
                                 const IntArray& axis,
                                 bool keep_dim,
                                 bool reduce_all) {
-  TensorDistAttr x_dist_attr = out_grad.dist_attr();
-  TensorDistAttr x_grad_dist_attr = out_grad.dist_attr();
+  TensorDistAttr out_grad_dist_attr = out_grad.dist_attr();
+  out_grad_dist_attr.clean_partial_status();
+  TensorDistAttr x_dist_attr = out_grad_dist_attr;
+  TensorDistAttr x_grad_dist_attr = out_grad_dist_attr;
 
   std::vector<int64_t> x_dim = phi::vectorize(x.dims());
   std::vector<int64_t> out_grad_dim = phi::vectorize(out_grad.dims());
@@ -241,7 +243,7 @@ SpmdInfo ReductionGradInferSpmd(const DistMetaTensor& x,
     x_grad_dist_attr.set_dims_mapping(dims_mapping);
   }
 
-  return {{x_dist_attr, out_grad.dist_attr()}, {x_grad_dist_attr}};
+  return {{x_dist_attr, out_grad_dist_attr}, {x_grad_dist_attr}};
 }
 
 }  // namespace distributed
diff --git a/test/auto_parallel/semi_auto_parallel_for_matmul.py b/test/auto_parallel/semi_auto_parallel_for_matmul.py
@@ -111,7 +111,6 @@ def test_matmul_x_column_shard(self):
         np.testing.assert_equal(
             dist_out.dist_attr.dims_mapping, [-1, -1], verbose=True
         )
-        assert dist_out.dist_attr._is_partial() is False
         # verify x_grad local shape and dist attr
         np.testing.assert_equal(
             dist_x_grad._local_shape, [64, 16], verbose=True
diff --git a/test/auto_parallel/test_semi_auto_parallel_hybrid_strategy.py b/test/auto_parallel/test_semi_auto_parallel_hybrid_strategy.py
@@ -19,7 +19,11 @@
 
 class TestSemiAutoParallelHybridStrategy(test_base.CommunicationTestDistBase):
     def setUp(self):
-        super().setUp(num_of_devices=2, timeout=120, nnode=2)
+        super().setUp(
+            num_of_devices=2,
+            timeout=120,
+            nnode=2,
+        )
         self._default_envs = {
             "dtype": "float32",
             "seed": "2023",

Original file line number	Diff line number	Diff line change
`@@ -111,7 +111,6 @@ def test_matmul_x_column_shard(self):`
`111`	`111`	`np.testing.assert_equal(`
`112`	`112`	`dist_out.dist_attr.dims_mapping, [-1, -1], verbose=True`
`113`	`113`	`)`
`114`		`- assert dist_out.dist_attr._is_partial() is False`
`115`	`114`	`# verify x_grad local shape and dist attr`
`116`	`115`	`np.testing.assert_equal(`
`117`	`116`	`dist_x_grad._local_shape, [64, 16], verbose=True`