[AutoParallel] Fix the bug in reshape_grad-like op when their inputs should be reshard (#67729)

pkuzyc · web-flow · commit 308d3c2336fd · 2024-08-27T14:39:29.000+08:00
* fix the bug that reshape_grad should compute local shape after infer_spmd

* remove comments
diff --git a/paddle/phi/api/generator/dist_api_gen.py b/paddle/phi/api/generator/dist_api_gen.py
@@ -1508,7 +1508,37 @@ def generate_prepare_data_code(self) -> str:
 
         return input_tensor_code, input_name_tensor_map
 
-    def generate_infer_meta_code(self) -> str:
+    def get_shape_type(self, attr_info):
+        shape_type = "int"
+        for name, info in attr_info.items():
+            if "IntArray" in info[0]:
+                shape_type = "int64_t"
+        return shape_type
+
+    def generate_infer_local_shape_code(self) -> str:
+        arg_name = self.infer_meta['local_shape']
+        assert arg_name in self.outputs['names'], (
+            f"Auto Parallel will calculate local_shape for {arg_name} "
+            f"in {self.api}, but {arg_name} is not found in its outputs."
+        )
+        # shape_type = self.attrs['attr_info'][shape_name][0]
+        # out_name = self.dist_output_args[0]
+        dist_out_name = self.dist_output_args[
+            self.outputs['names'].index(arg_name)
+        ]
+        shape_type = self.get_shape_type(self.attrs['attr_info'])
+        return CALCULATE_LOCAL_SHAPE_TEMPLATE.format(
+            out_name=dist_out_name,
+            out_dist_attr=(
+                "PADDLE_GET_CONST(phi::distributed::TensorDistAttr, spmd_info.second[0]);"
+                if self.infer_meta['spmd_rule']
+                else f"phi::distributed::TensorDistAttr(common::vectorize({dist_out_name}->dims()))"
+            ),
+            dtype=shape_type,
+            op_name=self.kernel['func'][0],
+        )
+
+    def generate_infer_meta_func_and_args_code(self) -> str:
         input_names = self.inputs['names']
         attr_names = self.attrs['names']
 
@@ -1583,6 +1613,20 @@ def generate_infer_meta_code(self) -> str:
                         f"{out_name} ? &meta_{out_name} : nullptr, "
                     )
         output_args_code = output_args_code[:-2]
+        return (
+            infer_meta_func_code,
+            input_args_code,
+            output_decl_code,
+            output_args_code,
+        )
+
+    def generate_infer_meta_code(self) -> str:
+        (
+            infer_meta_func_code,
+            input_args_code,
+            output_decl_code,
+            output_args_code,
+        ) = self.generate_infer_meta_func_and_args_code()
 
         infer_meta_code = ""
 
@@ -1602,25 +1646,8 @@ def generate_infer_meta_code(self) -> str:
 
         # TODO(GhostScreaming): kernel like reshape need calculate local_shape
         if self.infer_meta['local_shape'] is not None:
-            shape_name = self.infer_meta['local_shape']
-            assert (
-                shape_name in self.attrs['names']
-            ), f"Auto Parallel will calculate local_shape {shape_name} for"
-            "operator {self.kernel['func'][0]}, but {shape_name} is not"
-            "found in its attributes."
-            shape_type = self.attrs['attr_info'][shape_name][0]
-            out_name = self.dist_output_args[0]
-            infer_meta_code += CALCULATE_LOCAL_SHAPE_TEMPLATE.format(
-                out_name=out_name,
-                out_dist_attr=(
-                    "PADDLE_GET_CONST(phi::distributed::TensorDistAttr, spmd_info.second[0]);"
-                    if self.infer_meta['spmd_rule']
-                    else f"phi::distributed::TensorDistAttr(common::vectorize({out_name}->dims()))"
-                ),
-                dtype="int64_t" if shape_type == "IntArray" else "int",
-                op_name=self.kernel['func'][0],
-                shape_name=shape_name,
-            )
+            infer_meta_code += self.generate_infer_local_shape_code()
+
         infer_meta_code = infer_meta_code + INFER_META_TEMPLATE.format(
             infer_meta_func_code, input_args_code, output_args_code
         )
diff --git a/paddle/phi/api/generator/dist_bw_api_gen.py b/paddle/phi/api/generator/dist_bw_api_gen.py
@@ -14,6 +14,7 @@
 
 import argparse
 
+import dist_api_gen
 import yaml
 from backward_api_gen import BackwardAPI
 from dist_api_gen import DistForwardAPI
@@ -194,10 +195,14 @@
 NONEED_TO_RESHARD_OUTPUT_TEMPLATE = """
     // API `{}` does not need to reshard output."""
 
+SET_LOCAL_SHAPE_TEMPLATE = """
+      {meta_tensor}.set_dims(phi::make_ddim(local_shape));"""
+
 
 class DistBackwardAPI(DistForwardAPI, BackwardAPI):
     def __init__(self, backward_item_yaml):
         BackwardAPI.__init__(self, backward_item_yaml)
+        self.forward_config = backward_item_yaml['forward']
         self.init_dist_api_members()
 
     # override DistForwardAPI's method
@@ -283,6 +288,72 @@ def generate_output_creation_code(self) -> str:
 
         return output_creation_code
 
+    def generate_bw_infer_local_shape_code(self) -> str:
+        arg_name = self.infer_meta['local_shape']
+        assert arg_name in self.outputs['names'], (
+            f"Auto Parallel will calculate local_shape for {arg_name} "
+            f"in {self.api}, but {arg_name} is not found in its outputs."
+        )
+
+        _, fw_inputs, fw_attrs, fw_outputs = self.parse_forward_config(
+            self.forward_config
+        )
+        # shape_type = self.attrs['attr_info'][shape_name][0]
+        # out_name = self.dist_output_args[0]
+        dist_out_name = self.dist_output_args[
+            self.outputs['names'].index(arg_name)
+        ]
+        shape_type = self.get_shape_type(fw_attrs['attr_info'])
+        return dist_api_gen.CALCULATE_LOCAL_SHAPE_TEMPLATE.format(
+            out_name=dist_out_name,
+            out_dist_attr=(
+                "PADDLE_GET_CONST(phi::distributed::TensorDistAttr, spmd_info.second[0]);"
+                if self.infer_meta['spmd_rule']
+                else f"phi::distributed::TensorDistAttr(common::vectorize({dist_out_name}->dims()))"
+            ),
+            dtype=shape_type,
+            op_name=self.kernel['func'][0],
+        )
+
+    def generate_infer_meta_code(self) -> str:
+        (
+            infer_meta_func_code,
+            input_args_code,
+            output_decl_code,
+            output_args_code,
+        ) = self.generate_infer_meta_func_and_args_code()
+
+        infer_meta_code = ""
+
+        if self.infer_meta['global_shape'] is not None:
+            for i, out_name in enumerate(self.outputs['names']):
+                if out_name == self.infer_meta[
+                    'global_shape'
+                ] and self.need_to_generate_code_for_inplace_impl(i):
+                    infer_meta_code += dist_api_gen.SET_DIMS_TEMPLATE.format(
+                        dst=self.dist_output_args[i],
+                        src=(
+                            self.dist_output_args[i] + '_tmp'
+                            if i > 0
+                            else self.dist_output_args[i]
+                        ),
+                    )
+
+        infer_meta_code = (
+            infer_meta_code
+            + dist_api_gen.INFER_META_TEMPLATE.format(
+                infer_meta_func_code, input_args_code, output_args_code
+            )
+        )
+        # TODO(GhostScreaming): kernel like reshape need calculate local_shape
+        if self.infer_meta['local_shape'] is not None:
+            infer_meta_code += self.generate_bw_infer_local_shape_code()
+            infer_meta_code += SET_LOCAL_SHAPE_TEMPLATE.format(
+                meta_tensor="meta_" + self.dense_output_args[0]
+            )
+
+        return output_decl_code + infer_meta_code
+
     # override DistForwardAPI's method
     def generate_return_code(self) -> str:
         return "return;"
diff --git a/paddle/phi/infermeta/spmd_rules/reshape.cc b/paddle/phi/infermeta/spmd_rules/reshape.cc
@@ -336,12 +336,6 @@ SpmdInfo ReshapeGradInferSpmd(const DistMetaTensor& x,
   const auto& x_dist_dst = PADDLE_GET_CONST(TensorDistAttr, tmp.first[0]);
   const auto& out_grad_dist_dst =
       PADDLE_GET_CONST(TensorDistAttr, tmp.second[0]);
-  PADDLE_ENFORCE_EQ(
-      x_dist_tmp.dims_mapping(),
-      x_dist_dst.dims_mapping(),
-      common::errors::InvalidArgument("x should not be re shared: [%s] => [%s]",
-                                      x_dist_tmp.to_string(),
-                                      x_dist_dst.to_string()));
   return {{x_dist_tmp, out_grad_dist_dst}, {x_dist_dst}};
 }
 
diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml
@@ -2561,6 +2561,7 @@
     func : GradSameWithXInferMeta
     param : [x, out_grad]
     spmd_rule: ReshapeGradInferSpmd
+    local_shape : x_grad
   kernel :
     func : reshape_grad
     param : [x, out_grad]
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -1643,7 +1643,7 @@
   output : Tensor(out)
   infer_meta :
     func : ExpandAsInferMeta
-    local_shape: target_shape
+    local_shape: out
   kernel :
     func : expand_as
     data_type : x
@@ -3921,7 +3921,7 @@
   infer_meta :
     func : ReshapeInferMeta
     spmd_rule : ReshapeInferSpmd
-    local_shape: shape
+    local_shape: out
     global_shape: out
   kernel :
     func : reshape
