[Auto Paralle] update fused_gemm_epilogue_rule partial rule (#73494)

GITD245 · web-flow · commit 44634dd1d3f0 · 2025-06-21T15:14:06.000+08:00
diff --git a/paddle/phi/infermeta/spmd_rules/fused_gemm_epilogue.cc b/paddle/phi/infermeta/spmd_rules/fused_gemm_epilogue.cc
@@ -166,16 +166,15 @@ SpmdInfo FusedGemmEpilogueInferSpmdBase(const DistMetaTensor& x,
           "FusedGemmEpilogue, The ndim of bias should be 1, but got [%d].",
           bias_ndim));
 
-  VLOG(6) << "FusedGemmEpilogueSPMDRule InferForward Inputs: "
-          << "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) << "]; bias shape: ["
-          << str_join(ori_bias_shape) << "], bias_dims_mapping: ["
-          << str_join(bias_dims_mapping) << "]; trans_x: "
-          << "[" << (trans_x ? "true" : "false") << "]; "
-          << "trans_y: "
-          << "[" << (trans_y ? "true" : "false") << "]; ";
+  VLOG(4) << "FusedGemmEpilogueSPMDRule InferForward Inputs: ";
+  VLOG(4) << "X shape: [" << str_join(ori_x_shape) << "], x_dims_mapping: ["
+          << str_join(x_dims_mapping) << "];";
+  VLOG(4) << "Y shape: [" << str_join(ori_y_shape) << "], y_dims_mapping: ["
+          << str_join(y_dims_mapping) << "];";
+  VLOG(4) << "bias shape: [" << str_join(ori_bias_shape)
+          << "], bias_dims_mapping: [" << str_join(bias_dims_mapping) << "];";
+  VLOG(4) << "trans_x: [" << (trans_x ? "true" : "false") << "]; "
+          << "trans_y: [" << (trans_y ? "true" : "false") << "]; ";
 
   // Step1: build Einsum Notation
   std::string x_axes;
@@ -244,13 +243,19 @@ SpmdInfo FusedGemmEpilogueInferSpmdBase(const DistMetaTensor& x,
   output_dist_attr_dst.set_partial_status(partial_on_dims);
 
   if (output_dist_attr_dst.is_partial()) {
-    VLOG(4) << "FusedGemmEpilogue not support partial output, force set output "
-               "to replicated.";
-    output_dist_attr_dst.clean_partial_status();
-    SetTensorDistAttrReplicated(&x_dist_attr_dst, x_ndim);
-    SetTensorDistAttrReplicated(&y_dist_attr_dst, y_ndim);
-    SetTensorDistAttrReplicated(&bias_dist_attr_dst, bias_ndim);
-    SetTensorDistAttrReplicated(&output_dist_attr_dst, out_axes.size());
+    bias_dist_attr_dst.set_partial_status(
+        output_dist_attr_dst.partial_status());
+    if (!IsPartialLegal(bias_dist_attr_dst) ||
+        !IsPartialLegal(output_dist_attr_dst)) {
+      VLOG(4) << "FusedGemmEpilogue partial output illegal, force set output "
+                 "to replicated.";
+      output_dist_attr_dst.clean_partial_status();
+      bias_dist_attr_dst.clean_partial_status();
+      SetTensorDistAttrReplicated(&x_dist_attr_dst, x_ndim);
+      SetTensorDistAttrReplicated(&y_dist_attr_dst, y_ndim);
+      SetTensorDistAttrReplicated(&bias_dist_attr_dst, bias_ndim);
+      SetTensorDistAttrReplicated(&output_dist_attr_dst, out_axes.size());
+    }
   }
   TensorDistAttr output_reserve_dist_attr_dst =
       CopyTensorDistAttrForOutput(output_dist_attr_dst);
@@ -262,7 +267,6 @@ SpmdInfo FusedGemmEpilogueInferSpmdBase(const DistMetaTensor& x,
   LogInputDistAttr(
       "Bias", ori_bias_shape, bias_dist_attr_src, bias_dist_attr_dst);
   LogOutputDistAttr("Output", output_dist_attr_dst);
-  VLOG(4) << std::endl;
 
   return {{x_dist_attr_dst, y_dist_attr_dst, bias_dist_attr_dst},
           {output_dist_attr_dst, output_reserve_dist_attr_dst}};
diff --git a/paddle/phi/infermeta/spmd_rules/utils.cc b/paddle/phi/infermeta/spmd_rules/utils.cc
@@ -221,6 +221,7 @@ bool IsPartialLegal(const TensorDistAttr& dist_attr) {
     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) {
+        VLOG(4) << "Partial on dim [" << dim << "] but this dim is sharded";
         return false;
       }
     }
diff --git a/test/auto_parallel/spmd_rules/test_fused_gemm_epilogue_rule.py b/test/auto_parallel/spmd_rules/test_fused_gemm_epilogue_rule.py
@@ -23,7 +23,7 @@
 from paddle.framework import core
 
 
-class TestMatmulSPMDRule(unittest.TestCase):
+class TestFusedGemmEpilogueSPMDRule(unittest.TestCase):
     def setUp(self):
         # After replaced all spmd rules by phi impl, we can recover the
         # api name to `get_spmd_rule`
@@ -52,7 +52,7 @@ def test_fused_gemm_epilogue_infer_forward(self):
         )
         self.bias_dist_tensor_spec.set_dims_mapping([-1])
 
-        # has partial,force to replicate test partial: mk[1, 0],kn[0, -1],bias[-1] --> mk[1, 0],kn[0, -1],bias[-1] = nm[1, -1]
+        # mk[1, 0],kn[0, -1],bias[-1] --> mk[1, 0],kn[0, -1],bias[-1]partial[0] = nm[1, -1]partial[0]
         self.x_dist_tensor_spec.set_dims_mapping([1, 0])
         self.y_dist_tensor_spec.set_dims_mapping([0, -1])
 
@@ -71,11 +71,15 @@ def test_fused_gemm_epilogue_infer_forward(self):
         self.assertEqual(len(inferred_input_dist_attrs), 3)
         self.assertEqual(len(inferred_output_dist_attrs), 2)
 
-        self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [-1, -1])
-        self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, -1])
+        self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [1, 0])
+        self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [0, -1])
         self.assertEqual(inferred_input_dist_attrs[2].dims_mapping, [-1])
-        self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [-1, -1])
-        self.assertEqual(inferred_output_dist_attrs[0]._is_partial(), False)
+        self.assertEqual(inferred_input_dist_attrs[2]._is_partial(), True)
+        self.assertEqual(inferred_input_dist_attrs[2]._partial_dims(), {0})
+
+        self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [1, -1])
+        self.assertEqual(inferred_output_dist_attrs[0]._is_partial(), True)
+        self.assertEqual(inferred_output_dist_attrs[0]._partial_dims(), {0})
 
         # test row parallel: mk[1, -1],kn[-1, -1],bias[-1] --> mk[1, -1],kn[-1, -1],bias[-1] = nm[1, -1]
         self.x_dist_tensor_spec.set_dims_mapping([1, -1])
@@ -119,7 +123,7 @@ def test_fused_gemm_epilogue_infer_forward(self):
         self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [-1, 0])
         self.assertEqual(inferred_output_dist_attrs[0]._is_partial(), False)
 
-        # has partial,force to replicate test partial with propagation: mk[1, 0],kn[-1,-1],bias[-1] --> mk[1, 0],kn[0, -1],bias[-1] = nm[1, -1]
+        # mk[1, 0],kn[-1,-1],bias[-1] --> mk[1, 0],kn[0, -1],bias[-1]partial[0] = nm[1, -1]partial[0]
         self.x_dist_tensor_spec.set_dims_mapping([1, 0])
         self.y_dist_tensor_spec.set_dims_mapping([-1, -1])
         self.bias_dist_tensor_spec.set_dims_mapping([-1])
@@ -135,13 +139,17 @@ def test_fused_gemm_epilogue_infer_forward(self):
         inferred_input_dist_attrs = result_dist_attrs[0]
         inferred_output_dist_attrs = result_dist_attrs[1]
 
-        self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [-1, -1])
-        self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, -1])
+        self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [1, 0])
+        self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [0, -1])
         self.assertEqual(inferred_input_dist_attrs[2].dims_mapping, [-1])
-        self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [-1, -1])
-        self.assertEqual(inferred_output_dist_attrs[0]._is_partial(), False)
+        self.assertEqual(inferred_input_dist_attrs[2]._is_partial(), True)
+        self.assertEqual(inferred_input_dist_attrs[2]._partial_dims(), {0})
 
-        # has partial,force to replicate mk[-1,-1],kn[1,0],bias[-1] --> mk[-1, 1],kn[1, 0],bias[0] = nm[-1, 0]
+        self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [1, -1])
+        self.assertEqual(inferred_output_dist_attrs[0]._is_partial(), True)
+        self.assertEqual(inferred_output_dist_attrs[0]._partial_dims(), {0})
+
+        # mk[-1,-1],kn[1,0],bias[-1] --> mk[-1, 1],kn[1, 0],bias[0]partial[1] = nm[-1, 0]partial[1]
         self.x_dist_tensor_spec.set_dims_mapping([-1, -1])
         self.y_dist_tensor_spec.set_dims_mapping([1, 0])
         self.bias_dist_tensor_spec.set_dims_mapping([-1])
@@ -157,11 +165,15 @@ def test_fused_gemm_epilogue_infer_forward(self):
         inferred_input_dist_attrs = result_dist_attrs[0]
         inferred_output_dist_attrs = result_dist_attrs[1]
 
-        self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [-1, -1])
-        self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, -1])
-        self.assertEqual(inferred_input_dist_attrs[2].dims_mapping, [-1])
-        self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [-1, -1])
-        self.assertEqual(inferred_output_dist_attrs[0]._is_partial(), False)
+        self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [-1, 1])
+        self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [1, 0])
+        self.assertEqual(inferred_input_dist_attrs[2].dims_mapping, [0])
+        self.assertEqual(inferred_input_dist_attrs[2]._is_partial(), True)
+        self.assertEqual(inferred_input_dist_attrs[2]._partial_dims(), {1})
+
+        self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [-1, 0])
+        self.assertEqual(inferred_output_dist_attrs[0]._is_partial(), True)
+        self.assertEqual(inferred_output_dist_attrs[0]._partial_dims(), {1})
 
         # abcmk[1, 0, -1, -1],kn[-1, -1],bias[-1] --> abcmk[1, 0, -1, -1],kn[-1, -1],bias[-1] = abcmn[1, 0, -1, -1]
         self.x_dist_tensor_spec.shape = [512, 48, 64, 32]
@@ -187,7 +199,7 @@ def test_fused_gemm_epilogue_infer_forward(self):
         )
         self.assertEqual(inferred_output_dist_attrs[0]._is_partial(), False)
 
-        # has partial,force to replicate abcmk[1, -1, -1, 0],kn[-1, -1],bias[-1] --> abcmk[1, -1, -1, 0],kn[0, -1],bias[-1] = abcmn[1,-1, -1, -1]
+        # abcmk[1, -1, -1, 0],kn[-1, -1],bias[-1] --> abcmk[1, -1, -1, 0],kn[0, -1],bias[-1]partial[0] = abcmn[1,-1, -1, -1]partial[0]
         self.x_dist_tensor_spec.set_dims_mapping([1, -1, -1, 0])
         self.y_dist_tensor_spec.set_dims_mapping([-1, -1])
 
@@ -202,14 +214,17 @@ def test_fused_gemm_epilogue_infer_forward(self):
         inferred_input_dist_attrs = result_dist_attrs[0]
         inferred_output_dist_attrs = result_dist_attrs[1]
         self.assertEqual(
-            inferred_input_dist_attrs[0].dims_mapping, [1, -1, -1, -1]
+            inferred_input_dist_attrs[0].dims_mapping, [1, -1, -1, 0]
         )
-        self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, -1])
+        self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [0, -1])
         self.assertEqual(inferred_input_dist_attrs[2].dims_mapping, [-1])
+        self.assertEqual(inferred_input_dist_attrs[2]._is_partial(), True)
+        self.assertEqual(inferred_input_dist_attrs[2]._partial_dims(), {0})
         self.assertEqual(
             inferred_output_dist_attrs[0].dims_mapping, [1, -1, -1, -1]
         )
-        self.assertEqual(inferred_output_dist_attrs[0]._is_partial(), False)
+        self.assertEqual(inferred_output_dist_attrs[0]._is_partial(), True)
+        self.assertEqual(inferred_output_dist_attrs[0]._partial_dims(), {0})
 
         # trans_x = True, abcmk[1, -1, -1, 0], kn[-1, -1],bias[-1] --> abcmk[1, -1, -1, 0],kn[-1, -1],bias[-1] = abcmn[1, -1, 0, -1]
         self.x_dist_tensor_spec.set_dims_mapping([1, -1, -1, 0])
@@ -230,12 +245,13 @@ def test_fused_gemm_epilogue_infer_forward(self):
         )
         self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, -1])
         self.assertEqual(inferred_input_dist_attrs[2].dims_mapping, [-1])
+        self.assertEqual(inferred_input_dist_attrs[2]._is_partial(), False)
         self.assertEqual(
             inferred_output_dist_attrs[0].dims_mapping, [1, -1, 0, -1]
         )
         self.assertEqual(inferred_output_dist_attrs[0]._is_partial(), False)
 
-        # has partial,force to replicate trans_y = True, abcmk[-1, -1, -1, -1], kn[1, 0],bias[-1] --> abcmk[-1, -1, -1, 0],kn[1, 0],bias[1] = abcmn[-1, -1, -1, 1] partial[0]: done
+        # trans_y = True, abcmk[-1, -1, -1, -1], kn[1, 0],bias[-1] --> abcmk[-1, -1, -1, 0],kn[1, 0],bias[1]partial[0] = abcmn[-1, -1, -1, 1]partial[0]
         self.x_dist_tensor_spec.set_dims_mapping([-1, -1, -1, -1])
         self.y_dist_tensor_spec.set_dims_mapping([1, 0])
 
@@ -250,16 +266,19 @@ def test_fused_gemm_epilogue_infer_forward(self):
         inferred_input_dist_attrs = result_dist_attrs[0]
         inferred_output_dist_attrs = result_dist_attrs[1]
         self.assertEqual(
-            inferred_input_dist_attrs[0].dims_mapping, [-1, -1, -1, -1]
+            inferred_input_dist_attrs[0].dims_mapping, [-1, -1, -1, 0]
         )
-        self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, -1])
-        self.assertEqual(inferred_input_dist_attrs[2].dims_mapping, [-1])
+        self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [1, 0])
+        self.assertEqual(inferred_input_dist_attrs[2].dims_mapping, [1])
+        self.assertEqual(inferred_input_dist_attrs[2]._is_partial(), True)
+        self.assertEqual(inferred_input_dist_attrs[2]._partial_dims(), {0})
         self.assertEqual(
-            inferred_output_dist_attrs[0].dims_mapping, [-1, -1, -1, -1]
+            inferred_output_dist_attrs[0].dims_mapping, [-1, -1, -1, 1]
         )
-        self.assertEqual(inferred_output_dist_attrs[0]._is_partial(), False)
+        self.assertEqual(inferred_output_dist_attrs[0]._is_partial(), True)
+        self.assertEqual(inferred_output_dist_attrs[0]._partial_dims(), {0})
 
-        # has partial,force to replicate trans_y = True, trans_x = True, abcmk[-1, -1, 0, 1], kn[1, 0],bias[-1] --> abcmk[-1, -1, 0, 1]],kn[-1, 0],bias[-1] = abcmn[-1, -1, 1, -1]
+        # trans_y = True, trans_x = True, abcmk[-1, -1, 0, 1], kn[1, 0],bias[-1] --> abcmk[-1, -1, 0, 1],kn[-1, 0],bias[-1]partial[0] = abcmn[-1, -1, 1, -1]partial[0]
         # multiple mesh dim shard same tensor axis
         self.x_dist_tensor_spec.set_dims_mapping([-1, -1, 0, 1])
         self.y_dist_tensor_spec.set_dims_mapping([1, 0])
@@ -275,16 +294,19 @@ def test_fused_gemm_epilogue_infer_forward(self):
         inferred_input_dist_attrs = result_dist_attrs[0]
         inferred_output_dist_attrs = result_dist_attrs[1]
         self.assertEqual(
-            inferred_input_dist_attrs[0].dims_mapping, [-1, -1, -1, -1]
+            inferred_input_dist_attrs[0].dims_mapping, [-1, -1, 0, 1]
         )
-        self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, -1])
+        self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, 0])
         self.assertEqual(inferred_input_dist_attrs[2].dims_mapping, [-1])
+        self.assertEqual(inferred_input_dist_attrs[2]._is_partial(), True)
+        self.assertEqual(inferred_input_dist_attrs[2]._partial_dims(), {0})
         self.assertEqual(
-            inferred_output_dist_attrs[0].dims_mapping, [-1, -1, -1, -1]
+            inferred_output_dist_attrs[0].dims_mapping, [-1, -1, 1, -1]
         )
-        self.assertEqual(inferred_output_dist_attrs[0]._is_partial(), False)
+        self.assertEqual(inferred_output_dist_attrs[0]._is_partial(), True)
+        self.assertEqual(inferred_output_dist_attrs[0]._partial_dims(), {0})
 
-        # has partial,force to replicate mk[-1,1],k[1],bias[1] --> mk[-1,-1],k[-1],bias[-1] = m[-1]
+        # mk[-1,1],k[1],bias[1] --> mk[-1,1],k[1],bias[-1]partial[1] = m[-1]partial[1]
         x_shape = [64, 32]
         y_shape = [32]
         bias_shape = [32]
@@ -316,13 +338,17 @@ def test_fused_gemm_epilogue_infer_forward(self):
 
         inferred_input_dist_attrs = result_dist_attrs[0]
         inferred_output_dist_attrs = result_dist_attrs[1]
-        self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [-1, -1])
-        self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1])
+        self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [-1, 1])
+        self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [1])
         self.assertEqual(inferred_input_dist_attrs[2].dims_mapping, [-1])
+        self.assertEqual(inferred_input_dist_attrs[2]._is_partial(), True)
+        self.assertEqual(inferred_input_dist_attrs[2]._partial_dims(), {1})
+
         self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [-1])
-        self.assertEqual(inferred_output_dist_attrs[0]._is_partial(), False)
+        self.assertEqual(inferred_output_dist_attrs[0]._is_partial(), True)
+        self.assertEqual(inferred_output_dist_attrs[0]._partial_dims(), {1})
 
-        # has partial,force to replicate k[-1],kn[-1,1],bias[-1] --> k[-1],kn[-1,1],bias[1] = n[1]
+        # k[-1],kn[-1,1],bias[-1] --> k[-1],kn[-1,1],bias[1] = n[1]
         x_shape = [32]
         y_shape = [32, 64]
         bias_shape = [64]
@@ -357,6 +383,7 @@ def test_fused_gemm_epilogue_infer_forward(self):
         self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [-1])
         self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, 1])
         self.assertEqual(inferred_input_dist_attrs[2].dims_mapping, [1])
+        self.assertEqual(inferred_input_dist_attrs[2]._is_partial(), False)
         self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [1])
         self.assertEqual(inferred_output_dist_attrs[0]._is_partial(), False)
 

Original file line number	Diff line number	Diff line change
`@@ -221,6 +221,7 @@ bool IsPartialLegal(const TensorDistAttr& dist_attr) {`
`221`	`221`	`const std::set<int64_t> partial_on_dims = dist_attr.partial_dims();`
`222`	`222`	`for (const int64_t& dim : dims_mapping) {`
`223`	`223`	`if (dim != -1 && partial_on_dims.count(dim) != 0) {`
	`224`	`+ VLOG(4) << "Partial on dim [" << dim << "] but this dim is sharded";`
`224`	`225`	`return false;`
`225`	`226`	`}`
`226`	`227`	`}`