[Auto Parallel] add ElementwiseBinaryWithPartialInferSpmd for multiply (#73408)

GITD245 · web-flow · commit 3c5a95c28ce2 · 2025-06-20T18:29:20.000+08:00
diff --git a/paddle/phi/infermeta/spmd_rules/elementwise.cc b/paddle/phi/infermeta/spmd_rules/elementwise.cc
@@ -301,6 +301,102 @@ SpmdInfo ElementwiseBinaryInferSpmd(const DistMetaTensor& x,
   return {{x_dist_attr_dst, y_dist_attr_dst}, {out_dist_attr}};
 }
 
+// NOTE(nieyuntao): This function is only for `multiply` right now to support
+// partial propagation
+SpmdInfo ElementwiseBinaryWithPartialInferSpmd(const DistMetaTensor& x,
+                                               const DistMetaTensor& y) {
+  // Step0: Verify Input Args Based on Elementwise Logic
+  auto x_shape = common::vectorize(x.dims());
+  int x_ndim = static_cast<int>(x_shape.size());
+  auto y_shape = common::vectorize(y.dims());
+  int y_ndim = static_cast<int>(y_shape.size());
+  TensorDistAttr x_dist_attr_src = x.dist_attr();
+  TensorDistAttr y_dist_attr_src = y.dist_attr();
+  std::vector<int64_t> x_dims_mapping = x_dist_attr_src.dims_mapping();
+  std::vector<int64_t> y_dims_mapping = y_dist_attr_src.dims_mapping();
+  PADDLE_ENFORCE_EQ(
+      x_ndim,
+      x_dims_mapping.size(),
+      common::errors::InvalidArgument(
+          "ElementwiseBinaryWithPartial, The Tensor X's rank [%d] and X's "
+          "dims_mapping size [%d] are not matched.",
+          x_ndim,
+          x_dims_mapping.size()));
+  PADDLE_ENFORCE_EQ(
+      y_ndim,
+      y_dims_mapping.size(),
+      common::errors::InvalidArgument(
+          "ElementwiseBinaryWithPartial, The Tensor Y's rank [%d] and Y's "
+          "dims_mapping size [%d] are not matched.",
+          y_ndim,
+          y_dims_mapping.size()));
+
+  // Step1: Build Einsum Notation
+  std::string x_axes, y_axes, out_axes;
+  GetBinaryNotations(x_shape, y_shape, &x_axes, &y_axes, &out_axes);
+
+  // Step2: Sharding Propagation
+  // Step2.1: Merge input shardings
+  std::unordered_map<std::string, int64_t> axis_to_dim_map =
+      ShardingMergeForTensors(
+          {{x_axes, x_dims_mapping}, {y_axes, y_dims_mapping}});
+
+  // Step2.2: Infer output dims mapping from merged input dims mapping
+  std::vector<int64_t> out_dims_mapping =
+      GetDimsMappingForAxes(out_axes, axis_to_dim_map);
+
+  // initialize output dist_attr's process_mesh, batch_dim and dynamic dims with
+  // input dist_attr.
+  TensorDistAttr out_dist_attr = CopyTensorDistAttrForOutput(x_dist_attr_src);
+  out_dist_attr.set_dims_mapping(out_dims_mapping);
+
+  // Step2.3: Update inputs' dims mapping with merged one.
+  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(
+      GetDimsMappingForAxes(y_axes, axis_to_dim_map));
+
+  // Step3: Handle partial
+  if (x_dist_attr_src.is_partial() || y_dist_attr_src.is_partial()) {
+    if (x_dist_attr_src.is_partial() && y_dist_attr_src.is_partial()) {
+      if (x_ndim >= y_ndim) {
+        x_dist_attr_dst.set_partial_status(x_dist_attr_src.partial_status());
+        out_dist_attr.set_partial_status(x_dist_attr_src.partial_status());
+      } else {
+        y_dist_attr_dst.set_partial_status(y_dist_attr_src.partial_status());
+        out_dist_attr.set_partial_status(y_dist_attr_src.partial_status());
+      }
+    } else if (x_dist_attr_src.is_partial()) {
+      x_dist_attr_dst.set_partial_status(x_dist_attr_src.partial_status());
+      out_dist_attr.set_partial_status(x_dist_attr_src.partial_status());
+    } else if (y_dist_attr_src.is_partial()) {
+      y_dist_attr_dst.set_partial_status(y_dist_attr_src.partial_status());
+      out_dist_attr.set_partial_status(y_dist_attr_src.partial_status());
+    }
+    if (!IsPartialLegal(x_dist_attr_dst) || !IsPartialLegal(y_dist_attr_dst) ||
+        !IsPartialLegal(out_dist_attr)) {
+      x_dist_attr_dst.clean_partial_status();
+      y_dist_attr_dst.clean_partial_status();
+      out_dist_attr.clean_partial_status();
+    }
+  }
+
+  VLOG(4) << "ElementwiseWithPartialSPMDRule InferForward:";
+  VLOG(4) << "Input0 shape: [" << str_join(x_shape) << "] "
+          << "src_dims_mapping: [" << str_join(x_dims_mapping) << "] "
+          << "dst_dims_mapping: [" << str_join(x_dist_attr_dst.dims_mapping())
+          << "]";
+  VLOG(4) << "Input1 shape: [" << str_join(y_shape) << "] "
+          << "src_dims_mapping: [" << str_join(y_dims_mapping) << "] "
+          << "dst_dims_mapping: [" << str_join(y_dist_attr_dst.dims_mapping())
+          << "]";
+  VLOG(4) << "Output dims_mapping: [" + str_join(out_dims_mapping) + "]\n\n";
+
+  return {{x_dist_attr_dst, y_dist_attr_dst}, {out_dist_attr}};
+}
+
 SpmdInfo ElementwiseBinaryInferSpmdReverse(const DistMetaTensor& x,
                                            const DistMetaTensor& y,
                                            const DistMetaTensor& out) {
diff --git a/paddle/phi/infermeta/spmd_rules/elementwise.h b/paddle/phi/infermeta/spmd_rules/elementwise.h
@@ -41,6 +41,9 @@ SpmdInfo ElementwiseUnaryGradInferSpmd(const DistMetaTensor& x,
 SpmdInfo ElementwiseBinaryInferSpmd(const DistMetaTensor& x,
                                     const DistMetaTensor& y);
 
+SpmdInfo ElementwiseBinaryWithPartialInferSpmd(const DistMetaTensor& x,
+                                               const DistMetaTensor& y);
+
 SpmdInfo ElementwiseBinaryInferSpmdReverse(const DistMetaTensor& x,
                                            const DistMetaTensor& y,
                                            const DistMetaTensor& out);
diff --git a/paddle/phi/infermeta/spmd_rules/rules.cc b/paddle/phi/infermeta/spmd_rules/rules.cc
@@ -401,7 +401,7 @@ PD_REGISTER_SPMD_RULE(
     PD_INFER_SPMD(phi::distributed::ElementwiseBinaryInferSpmdReverse));
 PD_REGISTER_SPMD_RULE(
     multiply,
-    PD_INFER_SPMD(phi::distributed::ElementwiseBinaryInferSpmd),
+    PD_INFER_SPMD(phi::distributed::ElementwiseBinaryWithPartialInferSpmd),
     PD_INFER_SPMD(phi::distributed::ElementwiseBinaryInferSpmdReverse));
 PD_REGISTER_SPMD_RULE(
     elementwise_mul,
diff --git a/paddle/phi/infermeta/spmd_rules/utils.cc b/paddle/phi/infermeta/spmd_rules/utils.cc
@@ -215,6 +215,21 @@ bool PlacementEqual(const std::shared_ptr<PlacementStatus>& a,
   return a_shard->get_axis() == b_shard->get_axis();
 }
 
+bool IsPartialLegal(const TensorDistAttr& dist_attr) {
+  if (dist_attr.is_partial()) {
+    const std::vector<int64_t> dims_mapping = dist_attr.dims_mapping();
+    const std::set<int64_t> partial_on_dims = dist_attr.partial_dims();
+    for (const int64_t& dim : dims_mapping) {
+      if (dim != -1 && partial_on_dims.count(dim) != 0) {
+        return false;
+      }
+    }
+    return true;
+  } else {
+    return true;
+  }
+}
+
 void AlignDimsSharding(std::vector<TensorDistAttr>* input_attrs_ptr,
                        const std::vector<std::vector<int64_t>>& tensor_shapes,
                        const std::vector<std::string>& axis_names,
diff --git a/paddle/phi/infermeta/spmd_rules/utils.h b/paddle/phi/infermeta/spmd_rules/utils.h
@@ -98,6 +98,8 @@ TensorDistAttr UnShardTensorDim(const TensorDistAttr& dist_attr, int dim);
 bool PlacementEqual(const std::shared_ptr<PlacementStatus>& a,
                     const std::shared_ptr<PlacementStatus>& b);
 
+bool IsPartialLegal(const TensorDistAttr& dist_attr);
+
 void AlignDimsSharding(std::vector<TensorDistAttr>* input_attrs,
                        const std::vector<std::vector<int64_t>>& tensor_shapes,
                        const std::vector<std::string>& axis_names,
diff --git a/paddle/phi/ops/yaml/inconsistent/dygraph_ops.yaml b/paddle/phi/ops/yaml/inconsistent/dygraph_ops.yaml
@@ -279,7 +279,7 @@
   output : Tensor
   infer_meta :
     func : ElementwiseInferMeta
-    spmd_rule : ElementwiseBinaryInferSpmd
+    spmd_rule : ElementwiseBinaryWithPartialInferSpmd
   kernel :
     func : multiply {dense, dense -> dense},
            multiply_sr {selected_rows, dense -> selected_rows}
diff --git a/test/auto_parallel/spmd_rules/test_elementwise_rule.py b/test/auto_parallel/spmd_rules/test_elementwise_rule.py
@@ -702,6 +702,99 @@ def test_single_mesh_dim_broadcast_less_than(self):
         self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [0])
         self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [-1, 0])
 
+    def test_multi_mesh_dim_with_partial(self):
+        process_mesh = auto.ProcessMesh([[0, 1, 2], [3, 4, 5]])
+        self.multiply_rule = core.get_phi_spmd_rule("multiply")
+        x_shape = [96, 24, 48]
+        y_shape = [24, 48]
+
+        # [0, -1, -1] partial_dim=[1], [-1, 1] --> [0, -1, 1] partial_dim=[], [-1, 1], [0, -1, 1] partial_dim=[]
+        x_tensor_dist_attr = TensorDistAttr()
+        x_tensor_dist_attr.process_mesh = process_mesh
+        x_tensor_dist_attr.dims_mapping = [0, -1, -1]
+        x_tensor_dist_attr._set_partial_dims([1])
+        x_dist_tensor_spec = DistTensorSpec(x_shape, x_tensor_dist_attr)
+
+        y_tensor_dist_attr = TensorDistAttr()
+        y_tensor_dist_attr.process_mesh = process_mesh
+        y_tensor_dist_attr.dims_mapping = [-1, 1]
+        y_dist_tensor_spec = DistTensorSpec(y_shape, y_tensor_dist_attr)
+
+        resulted_dist_attrs = self.multiply_rule.infer_forward(
+            x_dist_tensor_spec, y_dist_tensor_spec
+        )
+        inferred_input_dist_attrs = resulted_dist_attrs[0]
+        inferred_output_dist_attrs = resulted_dist_attrs[1]
+
+        self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1, 1])
+        self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, 1])
+        self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, -1, 1])
+
+        self.assertEqual(inferred_input_dist_attrs[0]._is_partial(), False)
+        self.assertEqual(inferred_input_dist_attrs[1]._is_partial(), False)
+        self.assertEqual(inferred_output_dist_attrs[0]._is_partial(), False)
+
+        # [-1, -1, -1], [-1, 1] partial_dim=[0] --> [-1, -1, 1], [-1, 1] partial_dim=[0], [-1, -1, 1] partial_dim=[0]
+        x_tensor_dist_attr = TensorDistAttr()
+        x_tensor_dist_attr.process_mesh = process_mesh
+        x_tensor_dist_attr.dims_mapping = [-1, -1, -1]
+        x_dist_tensor_spec = DistTensorSpec(x_shape, x_tensor_dist_attr)
+
+        y_tensor_dist_attr = TensorDistAttr()
+        y_tensor_dist_attr.process_mesh = process_mesh
+        y_tensor_dist_attr.dims_mapping = [-1, 1]
+        y_tensor_dist_attr._set_partial_dims([0])
+        y_dist_tensor_spec = DistTensorSpec(y_shape, y_tensor_dist_attr)
+
+        resulted_dist_attrs = self.multiply_rule.infer_forward(
+            x_dist_tensor_spec, y_dist_tensor_spec
+        )
+        inferred_input_dist_attrs = resulted_dist_attrs[0]
+        inferred_output_dist_attrs = resulted_dist_attrs[1]
+
+        self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [-1, -1, 1])
+        self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, 1])
+        self.assertEqual(
+            inferred_output_dist_attrs[0].dims_mapping, [-1, -1, 1]
+        )
+
+        self.assertEqual(inferred_input_dist_attrs[0]._is_partial(), False)
+        self.assertEqual(inferred_input_dist_attrs[1]._is_partial(), True)
+        self.assertEqual(inferred_input_dist_attrs[1]._partial_dims(), {0})
+        self.assertEqual(inferred_output_dist_attrs[0]._is_partial(), True)
+        self.assertEqual(inferred_output_dist_attrs[0]._partial_dims(), {0})
+
+        # [-1, 0, -1] partial_dim=[1], [-1, 0] partial_dim=[1] --> [-1, 0, -1] partial_dim=[1], [0, -1] partial_dim=[], [-1, 0, -1] partial_dim=[1]
+        x_tensor_dist_attr = TensorDistAttr()
+        x_tensor_dist_attr.process_mesh = process_mesh
+        x_tensor_dist_attr.dims_mapping = [-1, 0, -1]
+        x_tensor_dist_attr._set_partial_dims([1])
+        x_dist_tensor_spec = DistTensorSpec(x_shape, x_tensor_dist_attr)
+
+        y_tensor_dist_attr = TensorDistAttr()
+        y_tensor_dist_attr.process_mesh = process_mesh
+        y_tensor_dist_attr.dims_mapping = [-1, 0]
+        y_tensor_dist_attr._set_partial_dims([1])
+        y_dist_tensor_spec = DistTensorSpec(y_shape, y_tensor_dist_attr)
+
+        resulted_dist_attrs = self.multiply_rule.infer_forward(
+            x_dist_tensor_spec, y_dist_tensor_spec
+        )
+        inferred_input_dist_attrs = resulted_dist_attrs[0]
+        inferred_output_dist_attrs = resulted_dist_attrs[1]
+
+        self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [-1, 0, -1])
+        self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [0, -1])
+        self.assertEqual(
+            inferred_output_dist_attrs[0].dims_mapping, [-1, 0, -1]
+        )
+
+        self.assertEqual(inferred_input_dist_attrs[0]._is_partial(), True)
+        self.assertEqual(inferred_input_dist_attrs[0]._partial_dims(), {1})
+        self.assertEqual(inferred_input_dist_attrs[1]._is_partial(), False)
+        self.assertEqual(inferred_output_dist_attrs[0]._is_partial(), True)
+        self.assertEqual(inferred_output_dist_attrs[0]._partial_dims(), {1})
+
 
 if __name__ == "__main__":
     unittest.main()
