[fluid_ops]dist_default modify c_allreduce_sum (#72111)

co63oc · web-flow · commit fbcbe532f8d4 · 2025-04-09T13:21:55.000+08:00
diff --git a/python/paddle/distributed/auto_parallel/static/operators/common.py b/python/paddle/distributed/auto_parallel/static/operators/common.py
@@ -664,12 +664,18 @@ def is_data_parallel_reduce_op(op):
         "c_allreduce_sum",
         "c_allreduce_avg",
     ]
+    is_all_reduce_op = op.type == "all_reduce" and op.desc.attr(
+        "reduce_type"
+    ) in [
+        dist.ReduceOp.SUM,
+        dist.ReduceOp.AVG,
+    ]
     is_reduce_op = op.type == "reduce" and op.desc.attr("reduce_type") in [
         dist.ReduceOp.SUM,
         dist.ReduceOp.AVG,
     ]
     return (
-        (is_allreduce_op or is_reduce_op)
+        (is_allreduce_op or is_all_reduce_op or is_reduce_op)
         and op.desc.has_attr("op_namescope")
         and ParallelMode.DataParallel in op.desc.attr("op_namescope")
     )
@@ -686,7 +692,13 @@ def is_amp_flag_sync_op(op):
 
 def is_global_norm_sync_op(op):
     return (
-        op.type == "c_allreduce_sum"
+        (
+            op.type == "c_allreduce_sum"
+            or (
+                op.type == "all_reduce"
+                and op.desc.attr("reduce_type") == dist.ReduceOp.SUM
+            )
+        )
         and op.desc.has_attr("op_namescope")
         and SyncMode.GlobalNormSync in op.desc.attr("op_namescope")
     )
diff --git a/python/paddle/distributed/auto_parallel/static/operators/dist_default.py b/python/paddle/distributed/auto_parallel/static/operators/dist_default.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License
 
-
+import paddle
 from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY, OpRole
 
 from ..completion import contains_spmd_rule, get_phi_spmd_rule
@@ -67,12 +67,12 @@ def prim_operator_data_parallel_functor(ctx, src_op):
         sync_group = new_process_group(ctx.data_parallel_group)
 
         allreduce_op = main_block.append_op(
-            type='c_allreduce_sum',
-            inputs={'X': [var_name]},
-            outputs={'Out': [var_name]},
+            type='all_reduce',
+            inputs={'x': [var_name]},
+            outputs={'out': [var_name]},
             attrs={
                 'ring_id': sync_group.id,
-                'use_calc_stream': True,
+                'reduce_type': paddle.distributed.ReduceOp.SUM,
                 OP_ROLE_KEY: OpRole.Backward,
             },
         )
diff --git a/test/deprecated/auto_parallel/test_while_op_partition_deprecated.py b/test/deprecated/auto_parallel/test_while_op_partition_deprecated.py
@@ -393,8 +393,13 @@ def test_partitioner(self):
         sub_block_ops = dist_main_prog.blocks[1].ops
         sub_block_ops = [op.type for op in sub_block_ops]
 
-        self.assertTrue("c_allreduce_sum" in global_block_ops)
-        self.assertTrue("c_allreduce_sum" in sub_block_ops)
+        self.assertTrue(
+            "c_allreduce_sum" in global_block_ops
+            or "all_reduce" in global_block_ops
+        )
+        self.assertTrue(
+            "c_allreduce_sum" in sub_block_ops or "all_reduce" in sub_block_ops
+        )
 
         # test fill_constant_batch_size_like
         self.assertIsNotNone(fill_op)
diff --git a/test/deprecated/legacy_test/test_auto_parallel_partitioner_gpt_deprecated.py b/test/deprecated/legacy_test/test_auto_parallel_partitioner_gpt_deprecated.py
@@ -962,7 +962,14 @@ def test_gpt_dp_mp(self):
                 op.desc.output_arg_names()[0].split("@")[0]
                 for op in auto_parallel_main_prog.global_block().ops
                 if (
-                    op.type == "c_allreduce_sum"
+                    (
+                        op.type == "c_allreduce_sum"
+                        or (
+                            op.type == "all_reduce"
+                            and op.attr('reduce_type')
+                            == paddle.distributed.ReduceOp.SUM
+                        )
+                    )
                     and op.attr('op_role') == 1
                     and op.desc.attr("ring_id") == mp_ring_id
                 )
@@ -973,7 +980,14 @@ def test_gpt_dp_mp(self):
                 op.desc.output_arg_names()[0].split("@")[0]
                 for op in auto_parallel_main_prog.global_block().ops
                 if (
-                    op.type == "c_allreduce_sum"
+                    (
+                        op.type == "c_allreduce_sum"
+                        or (
+                            op.type == "all_reduce"
+                            and op.attr('reduce_type')
+                            == paddle.distributed.ReduceOp.SUM
+                        )
+                    )
                     and op.desc.attr("ring_id") == dp_ring_id
                 )
             ]
