fix reshard of empty tensor and refine local shape calculation (#63798)

zhiqiu · web-flow · commit d67a83542b10 · 2024-04-25T11:10:08.000+08:00
* refine default dist_attr of output and fix reshard of empty tensor

* fix adamax_
diff --git a/paddle/phi/api/yaml/generator/dist_api_gen.py b/paddle/phi/api/yaml/generator/dist_api_gen.py
@@ -49,6 +49,7 @@
     // 1. InferSpmd (Infer DistAttr of Inputs&Outputs){}
     // 2. Create API Output & Prepare Dist and Dense Output{}
     // 3. Infer DistTensor's Global Shape{}\n
+
     if (rank_is_in_current_mesh) {{
       // 4. Select Kernel{}
       // 5. Reshard Input{}\n
@@ -489,37 +490,12 @@
 CALCULATE_LOCAL_SHAPE_TEMPLATE = """
 
       // The dist_input_x is a dist tensor, the dims() func return the global dims.
-      auto x_shape = dist_input_x->dims();
-      auto x_numel = dist_input_x->numel();
-      bool visit_negative = false;
-      auto global_shape = {shape};
+      auto out_shape = {out_name}->dims();
       std::vector<{dtype}> local_shape;
-      for (size_t i = 0; i < global_shape.size(); i++) {{
-        auto& out_dist_attr = PADDLE_GET_CONST(phi::distributed::TensorDistAttr, spmd_info.second[0]);
+      const auto& out_dist_attr = {out_dist_attr};
+      for (int i = 0; i < out_shape.size(); i++) {{
         if (out_dist_attr.dims_mapping()[i] >= 0) {{
-          {dtype} shape_i = global_shape[i];
-          if (shape_i == 0) {{
-            shape_i = x_shape[i];
-          }} else if (shape_i == -1) {{
-            PADDLE_ENFORCE(not visit_negative,
-                           phi::errors::InvalidArgument(
-                               "{op_name} can only have one -1 in the {shape_name}."));
-            visit_negative = true;
-            int64_t non_negative_product = 1;
-            for (size_t j = 0; j < global_shape.size(); j++) {{
-              if (i == j) {{
-                continue;
-              }}
-              int64_t tmp_j = global_shape[j];
-              if (tmp_j == 0) {{
-                tmp_j = x_shape[j];
-              }}
-              non_negative_product *= tmp_j;
-            }}
-            PADDLE_ENFORCE(x_numel % non_negative_product == 0,
-                           phi::errors::InvalidArgument("Cannot infer real shape for -1."));
-            shape_i = x_numel / non_negative_product;
-          }}
+          {dtype} shape_i = out_shape[i];
           int64_t dim = out_dist_attr.dims_mapping()[i];
           int64_t mesh_dim = out_dist_attr.process_mesh().shape()[dim];
           // TODO: Support aliquant condition.
@@ -530,7 +506,7 @@
                     "and shard mesh dims is %lld.", i, shape_i, mesh_dim));
           local_shape.push_back(shape_i / mesh_dim);
         }} else {{
-          local_shape.push_back({shape}[i]);
+          local_shape.push_back(out_shape[i]);
         }}
       }}
 """
@@ -1587,11 +1563,12 @@ def generate_infer_meta_code(self) -> str:
             "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(
-                shape=f"{shape_name}.GetData()"
-                if shape_type == "IntArray"
-                else f"{shape_name}",
+                out_name=out_name,
+                out_dist_attr="paddle::get<0>(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,
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/p_to_s_reshard_function.cc b/paddle/phi/core/distributed/auto_parallel/reshard/p_to_s_reshard_function.cc
@@ -64,7 +64,7 @@ void ReshardPToSWithPadding(DeviceContext* dev_ctx,
     std::swap(axis[0], axis[split_axis]);
     RESHARD_FUNCTOR(dev_ctx, Transpose, dtype, in, axis, &in_reduce_scatter);
   } else {
-    in_reduce_scatter.ShareDataWith(in);
+    in_reduce_scatter.ShareDataNoCheckWith(in);
   }
 
   DenseTensor out_reduce_scatter;
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/s_to_s_reshard_function.cc b/paddle/phi/core/distributed/auto_parallel/reshard/s_to_s_reshard_function.cc
@@ -94,7 +94,7 @@ void SToSReshardFunction::Eval(phi::DeviceContext* dev_ctx,
     RESHARD_FUNCTOR(
         dev_ctx, Reshape, dtype, out_transpose, pre_shape_vec, &in_all_to_all);
   } else {
-    in_all_to_all.ShareDataWith(in.value());
+    in_all_to_all.ShareDataNoCheckWith(in.value());
   }
 
   // 2. use all to all to switch data to other ranks
diff --git a/test/auto_parallel/semi_auto_parallel_for_fused_rope.py b/test/auto_parallel/semi_auto_parallel_for_fused_rope.py
@@ -61,10 +61,10 @@ def test_only_q_input(self):
         dist_q = dist.shard_tensor(q, self._mesh, dist.Shard(0))
         dist_q.stop_gradient = False
         dist_out_q, _, _ = fused_rotary_position_embedding(
-            q=dist_q, use_neox_rotary_style=False
+            q=dist_q, use_neox_rotary_style=True
         )
         out_q, _, _ = fused_rotary_position_embedding(
-            q, use_neox_rotary_style=False
+            q, use_neox_rotary_style=True
         )
         self.check_tensor_eq(out_q, dist_out_q)
         self.check_placements(dist_out_q, [dist.Shard(0)])
@@ -85,10 +85,10 @@ def test_only_q_input_time_major(self):
         dist_q.stop_gradient = False
 
         dist_out_q, _, _ = fused_rotary_position_embedding(
-            q=dist_q, use_neox_rotary_style=False, time_major=True
+            q=dist_q, use_neox_rotary_style=True, time_major=True
         )
         out_q, _, _ = fused_rotary_position_embedding(
-            q, use_neox_rotary_style=False, time_major=True
+            q, use_neox_rotary_style=True, time_major=True
         )
         self.check_tensor_eq(out_q, dist_out_q)
         # NOTE: fused_rope have not supported shard on seq_len, so reshard to dist.Replicate
diff --git a/test/collective/test_communication_api_base.py b/test/collective/test_communication_api_base.py
@@ -56,9 +56,9 @@ def run_test_case(self, script_file, user_defined_envs=None):
             runtime_envs.update(user_defined_envs)
         runtime_envs["CUDA_VISIBLE_DEVICES"] = self._devices
         if self._nnode > 1:
-            start_command = f"{self._python_interp} -u -m paddle.distributed.launch --nnode={self._nnode} --master=127.0.0.1:{self._find_free_port()} --log_dir {self._log_dir.name} --devices {self._devices} {script_file}"
+            start_command = f"{self._python_interp} -u -m paddle.distributed.launch --nnode={self._nnode} --master=127.0.0.1:{self._find_free_port()} --log_dir ./log --devices {self._devices} {script_file}"
         else:
-            start_command = f"{self._python_interp} -u -m paddle.distributed.launch --log_dir {self._log_dir.name} --devices {self._devices} {script_file}"
+            start_command = f"{self._python_interp} -u -m paddle.distributed.launch --log_dir ./log --devices {self._devices} {script_file}"
         start_command_list = start_command.strip().split()
 
         if self._nnode > 1:

Original file line number	Diff line number	Diff line change
`@@ -64,7 +64,7 @@ void ReshardPToSWithPadding(DeviceContext* dev_ctx,`
`64`	`64`	`std::swap(axis[0], axis[split_axis]);`
`65`	`65`	`RESHARD_FUNCTOR(dev_ctx, Transpose, dtype, in, axis, &in_reduce_scatter);`
`66`	`66`	`} else {`
`67`		`- in_reduce_scatter.ShareDataWith(in);`
	`67`	`+ in_reduce_scatter.ShareDataNoCheckWith(in);`
`68`	`68`	`}`
`69`	`69`
`70`	`70`	`DenseTensor out_reduce_scatter;`
Original file line number	Diff line number	Diff line change
`@@ -94,7 +94,7 @@ void SToSReshardFunction::Eval(phi::DeviceContext* dev_ctx,`
`94`	`94`	`RESHARD_FUNCTOR(`
`95`	`95`	`dev_ctx, Reshape, dtype, out_transpose, pre_shape_vec, &in_all_to_all);`
`96`	`96`	`} else {`
`97`		`- in_all_to_all.ShareDataWith(in.value());`
	`97`	`+ in_all_to_all.ShareDataNoCheckWith(in.value());`
`98`	`98`	`}`
`99`	`99`
`100`	`100`	`// 2. use all to all to switch data to other ranks`
Original file line number	Diff line number	Diff line change
`@@ -61,10 +61,10 @@ def test_only_q_input(self):`
`61`	`61`	`dist_q = dist.shard_tensor(q, self._mesh, dist.Shard(0))`
`62`	`62`	`dist_q.stop_gradient = False`
`63`	`63`	`dist_out_q, _, _ = fused_rotary_position_embedding(`
`64`		`- q=dist_q, use_neox_rotary_style=False`
	`64`	`+ q=dist_q, use_neox_rotary_style=True`
`65`	`65`	`)`
`66`	`66`	`out_q, _, _ = fused_rotary_position_embedding(`
`67`		`- q, use_neox_rotary_style=False`
	`67`	`+ q, use_neox_rotary_style=True`
`68`	`68`	`)`
`69`	`69`	`self.check_tensor_eq(out_q, dist_out_q)`
`70`	`70`	`self.check_placements(dist_out_q, [dist.Shard(0)])`
`@@ -85,10 +85,10 @@ def test_only_q_input_time_major(self):`
`85`	`85`	`dist_q.stop_gradient = False`
`86`	`86`
`87`	`87`	`dist_out_q, _, _ = fused_rotary_position_embedding(`
`88`		`- q=dist_q, use_neox_rotary_style=False, time_major=True`
	`88`	`+ q=dist_q, use_neox_rotary_style=True, time_major=True`
`89`	`89`	`)`
`90`	`90`	`out_q, _, _ = fused_rotary_position_embedding(`
`91`		`- q, use_neox_rotary_style=False, time_major=True`
	`91`	`+ q, use_neox_rotary_style=True, time_major=True`
`92`	`92`	`)`
`93`	`93`	`self.check_tensor_eq(out_q, dist_out_q)`
`94`	`94`	`# NOTE: fused_rope have not supported shard on seq_len, so reshard to dist.Replicate`