【SPMD】fix matmul spmd rule for wrong shape (#67059)

liym27 · web-flow · commit 8db45642b2f3 · 2024-08-08T11:51:48.000+08:00
diff --git a/paddle/phi/infermeta/spmd_rules/layer_norm.cc b/paddle/phi/infermeta/spmd_rules/layer_norm.cc
@@ -26,26 +26,6 @@ namespace distributed {
 
 using phi::distributed::auto_parallel::str_join;
 
-void LogInputDistAttr(const std::string& name,
-                      const std::vector<int64_t>& shape,
-                      const TensorDistAttr& src_dist_attr,
-                      const TensorDistAttr& dst_dist_attr) {
-  VLOG(4) << name << " shape: [" << str_join(shape) << "] "
-          << "src_dims_mapping: [" << str_join(src_dist_attr.dims_mapping())
-          << "] "
-          << "dst_dims_mapping: [" << str_join(dst_dist_attr.dims_mapping())
-          << "] "
-          << "src_partial: " << src_dist_attr.partial_status_string()
-          << " dst_partial: " << dst_dist_attr.partial_status_string();
-}
-
-void LogOutputDistAttr(const std::string& name,
-                       const TensorDistAttr& dst_dist_attr) {
-  VLOG(4) << name << " dims mapping: ["
-          << str_join(dst_dist_attr.dims_mapping()) << "] "
-          << "partial: " << dst_dist_attr.partial_status_string();
-}
-
 SpmdInfo LayerNormInferSpmd(const DistMetaTensor& x,
                             const DistMetaTensor& scale,
                             const DistMetaTensor& bias,
diff --git a/paddle/phi/infermeta/spmd_rules/matmul.cc b/paddle/phi/infermeta/spmd_rules/matmul.cc
@@ -118,10 +118,10 @@ SpmdInfo MatmulInferSpmd(const DistMetaTensor& x,
                          bool trans_x,
                          bool trans_y) {
   // Step0: verify input args based on matmul logic
-  auto x_shape = common::vectorize(x.dims());
-  auto y_shape = common::vectorize(y.dims());
-  int x_ndim = static_cast<int>(x_shape.size());
-  int y_ndim = static_cast<int>(y_shape.size());
+  auto ori_x_shape = common::vectorize(x.dims());
+  auto ori_y_shape = common::vectorize(y.dims());
+  int x_ndim = static_cast<int>(ori_x_shape.size());
+  int y_ndim = static_cast<int>(ori_y_shape.size());
   const auto& x_dist_attr_src = x.dist_attr();
   const auto& y_dist_attr_src = y.dist_attr();
   std::vector<int64_t> x_dims_mapping = x_dist_attr_src.dims_mapping();
@@ -142,10 +142,10 @@ SpmdInfo MatmulInferSpmd(const DistMetaTensor& x,
                                       y_dims_mapping.size()));
 
   VLOG(6) << "MatmulSPMDRule InferForward Inputs: "
-          << "X shape: [" << str_join(x_shape) << "], x_dims_mapping: ["
-          << str_join(x_dims_mapping) << "]; Y shape: [" << str_join(y_shape)
-          << "], y_dims_mapping: [" << str_join(y_dims_mapping)
-          << "]; trans_x: "
+          << "X shape: [" << str_join(ori_x_shape) << "], x_dims_mapping: ["
+          << str_join(x_dims_mapping) << "]; Y shape: ["
+          << str_join(ori_y_shape) << "], y_dims_mapping: ["
+          << str_join(y_dims_mapping) << "]; trans_x: "
           << "[" << (trans_x ? "true" : "false") << "]; "
           << "trans_y: "
           << "[" << (trans_y ? "true" : "false") << "]; ";
@@ -191,6 +191,14 @@ SpmdInfo MatmulInferSpmd(const DistMetaTensor& x,
   output_dist_attr_dst.set_dims_mapping(out_dims_mapping);
 
   // Step2.3: Merge and get Inputs' New Dims Mapping.
+  auto x_shape = common::vectorize(x.dims());
+  auto y_shape = common::vectorize(y.dims());
+  if (trans_x) {
+    std::iter_swap(x_shape.end() - 2, x_shape.end() - 1);
+  }
+  if (trans_y) {
+    std::iter_swap(y_shape.end() - 2, y_shape.end() - 1);
+  }
   TensorDistAttr x_dist_attr_dst = GetMatmulInferedDistAttr(
       x_dist_attr_src, x_shape, x_axes, axis_to_dim_map, trans_x);
   TensorDistAttr y_dist_attr_dst = GetMatmulInferedDistAttr(
@@ -205,18 +213,11 @@ SpmdInfo MatmulInferSpmd(const DistMetaTensor& x,
   // Step2.3.2  handle input tensor partial (TODO)
   VLOG(4) << "MatmulSPMDRule InferForward: "
           << "Einsum notation: [" << x_axes << "," << y_axes << " --> "
-          << out_axes << "]. " << std::endl
-          << "X shape: [" << str_join(x_shape) << "], src_dims_mapping: ["
-          << str_join(x_dist_attr_src.dims_mapping())
-          << "], dst_dims_mapping: ["
-          << str_join(x_dist_attr_dst.dims_mapping()) << "]; Y shape: ["
-          << str_join(y_shape) << "], src_dims_mapping: ["
-          << str_join(y_dist_attr_src.dims_mapping())
-          << "], dst_dims_mapping: ["
-          << str_join(y_dist_attr_dst.dims_mapping())
-          << "]; Output dims_mapping: [" << str_join(out_dims_mapping)
-          << "], partial_on_dims: [" << str_join(partial_on_dims) << "]";
-
+          << out_axes << "]. " << std::endl;
+  LogInputDistAttr("X", ori_x_shape, x_dist_attr_src, x_dist_attr_dst);
+  LogInputDistAttr("Y", ori_y_shape, y_dist_attr_src, y_dist_attr_dst);
+  LogOutputDistAttr("Output", output_dist_attr_dst);
+  VLOG(4) << std::endl;
   return {{x_dist_attr_dst, y_dist_attr_dst}, {output_dist_attr_dst}};
 }
 
@@ -266,13 +267,11 @@ SpmdInfo MatmulInferSpmdReverse(const DistMetaTensor& x,
 
   VLOG(4) << "MatmulSPMDRule InferBackward: "
           << "Einsum notation: [" << x_axes << "," << y_axes << " --> "
-          << out_axes << "]. " << std::endl
-          << "Out shape: [" << str_join(out_shape) << "], src_dims_mapping: ["
-          << str_join(out_dims_mapping) << "], dst_dims_mapping: ["
-          << str_join(out_dims_mapping) << "]; Input X dims_mapping: ["
-          << str_join(x_dist_attr_dst.dims_mapping())
-          << "], Input Y dims_mapping:["
-          << str_join(y_dist_attr_dst.dims_mapping()) << "].";
+          << out_axes << "]. " << std::endl;
+  LogInputDistAttr("Out", out_shape, out_dist_attr_src, out_dist_attr_src);
+  LogOutputDistAttr("Input X", x_dist_attr_dst);
+  LogOutputDistAttr("Input Y", y_dist_attr_dst);
+  VLOG(4) << std::endl;
 
   return {{x_dist_attr_dst, y_dist_attr_dst}, {out_dist_attr_src}};
 }
diff --git a/paddle/phi/infermeta/spmd_rules/utils.h b/paddle/phi/infermeta/spmd_rules/utils.h
@@ -214,5 +214,19 @@ TensorDistAttr ReduceGradBroadCastDims(const TensorDistAttr& input,
 TensorDistAttr ReduceGradBroadCastDims(int64_t input_dims,
                                        const TensorDistAttr& grad);
 
+#define LogInputDistAttr(name, shape, src_dist_attr, dst_dist_attr)          \
+  VLOG(4) << name << " shape: [" << str_join(shape) << "] "                  \
+          << "src_dims_mapping: [" << str_join(src_dist_attr.dims_mapping()) \
+          << "] "                                                            \
+          << "dst_dims_mapping: [" << str_join(dst_dist_attr.dims_mapping()) \
+          << "] "                                                            \
+          << "src_partial: " << src_dist_attr.partial_status_string()        \
+          << " dst_partial: " << dst_dist_attr.partial_status_string();
+
+#define LogOutputDistAttr(name, dst_dist_attr)              \
+  VLOG(4) << name << " dims mapping: ["                     \
+          << str_join(dst_dist_attr.dims_mapping()) << "] " \
+          << "partial: " << dst_dist_attr.partial_status_string();
+
 }  // namespace distributed
 }  // namespace phi
