[fluid_ops] dist_matmul modify c_allreduce_sum (#72126)

co63oc · web-flow · commit f95c494a7372 · 2025-04-10T11:00:15.000+08:00
diff --git a/python/paddle/distributed/auto_parallel/static/operators/dist_matmul.py b/python/paddle/distributed/auto_parallel/static/operators/dist_matmul.py
@@ -14,7 +14,9 @@
 
 import copy
 
+import paddle
 from paddle.distributed.auto_parallel.static.cost.comm_op_cost import (
+    AllReduceOpCost,
     AllreduceSumOpCost,
     IdentityOpCost,
 )
@@ -439,12 +441,12 @@ def _right_operand_parameter_matmul_backward(ctx, *args, **kwargs):
         )
         group = new_process_group(group_ranks)
         c_allreduce_sum_op = main_block.append_op(
-            type='c_allreduce_sum',
-            inputs={'X': kwargs['X@GRAD']},
-            outputs={'Out': kwargs['X@GRAD']},
+            type='all_reduce',
+            inputs={'x': kwargs['X@GRAD']},
+            outputs={'out': kwargs['X@GRAD']},
             attrs={
                 'ring_id': group.id,
-                'use_calc_stream': True,
+                'reduce_type': paddle.distributed.ReduceOp.SUM,
                 'use_model_parallel': True,
                 OP_ROLE_KEY: OpRole.Backward,
             },
@@ -1334,18 +1336,18 @@ def calc_bwd_cost(self, dist_op, ctx, cluster):
 
         # calc comm op cost
         if has_x_grad:
-            attrs = {"use_calc_stream": True, "use_model_parallel": True}
+            attrs = {"reduce_type": paddle.distributed.ReduceOp.SUM}
             var_names = backward_op.output("X@GRAD")
             c_allreduce_sum_desc_mapping = build_comm_desc_from_dist_op(
-                "c_allreduce_sum",
+                "all_reduce",
                 dist_op,
                 ctx,
                 var_names,
                 attrs=attrs,
                 parallel_axis=parallel_axis,
             )
             comm_op_cost_list = build_comm_costs_from_descs(
-                AllreduceSumOpCost,
+                AllReduceOpCost,
                 ctx,
                 processes,
                 c_allreduce_sum_desc_mapping,
@@ -1631,11 +1633,11 @@ def calc_fwd_cost(self, dist_op, ctx, cluster):
         parallel_axis = dist_op.dist_attr.get_input_dims_mapping(
             serial_op.input("Y")[0]
         )[-2]
-        attrs = {"use_calc_stream": True, "use_model_parallel": True}
+        attrs = {"reduce_type": paddle.distributed.ReduceOp.SUM}
 
         var_names = serial_op.output("Out")
         c_allreduce_sum_desc_mapping = build_comm_desc_from_dist_op(
-            "c_allreduce_sum",
+            "all_reduce",
             dist_op,
             ctx,
             var_names,
@@ -1644,7 +1646,7 @@ def calc_fwd_cost(self, dist_op, ctx, cluster):
         )
 
         comm_op_cost_list = build_comm_costs_from_descs(
-            AllreduceSumOpCost,
+            AllReduceOpCost,
             ctx,
             processes,
             c_allreduce_sum_desc_mapping,
@@ -1781,13 +1783,12 @@ def forward(ctx, *args, **kwargs):
         )
 
         c_allreduce_sum_op = main_block.append_op(
-            type='c_allreduce_sum',
-            inputs={'X': Out_var},
-            outputs={'Out': Out_var},
+            type='all_reduce',
+            inputs={'x': Out_var},
+            outputs={'out': Out_var},
             attrs={
                 'ring_id': group.id,
-                'use_calc_stream': True,
-                'use_model_parallel': True,
+                'reduce_type': paddle.distributed.ReduceOp.SUM,
                 OP_ROLE_KEY: src_op.attr('op_role'),
             },
         )
diff --git a/python/paddle/distributed/passes/auto_parallel_fused_linear_promotion.py b/python/paddle/distributed/passes/auto_parallel_fused_linear_promotion.py
@@ -17,6 +17,7 @@
 
 import numpy as np
 
+import paddle
 from paddle.distributed.auto_parallel.static.utils import (
     is_optimize_op,
     is_recompute_op,
@@ -52,7 +53,7 @@
 FUSED_LINEAR_SOURCE_PATTERNS_LIST = [
     # amp_level == 'o2' or 'o3'
     {  # only MP
-        "forward": ["matmul_v2", "c_allreduce_sum", "elementwise_add"],
+        "forward": ["matmul_v2", "all_reduce", "elementwise_add"],
         "backward": ["elementwise_add_grad", "matmul_v2_grad"],
     },
     {  # MP + SP
@@ -67,7 +68,7 @@
         ],
     },
     {  # DP + MP
-        "forward": ["matmul_v2", "c_allreduce_sum", "elementwise_add"],
+        "forward": ["matmul_v2", "all_reduce", "elementwise_add"],
         "backward": [
             "elementwise_add_grad",
             "c_allreduce_sum",
@@ -90,7 +91,7 @@
     },
     # amp_level == 'o1'
     {
-        "forward": ["matmul_v2", "c_allreduce_sum", "cast", "elementwise_add"],
+        "forward": ["matmul_v2", "all_reduce", "cast", "elementwise_add"],
         "backward": ["elementwise_add_grad", "matmul_v2_grad"],
     },
     {
@@ -105,7 +106,7 @@
         ],
     },
     {
-        "forward": ["matmul_v2", "c_allreduce_sum", "cast", "elementwise_add"],
+        "forward": ["matmul_v2", "all_reduce", "cast", "elementwise_add"],
         "backward": [
             "elementwise_add_grad",
             "c_allreduce_sum",
@@ -541,7 +542,11 @@ def _transform_forward_segment(
                     origin_comm_op.input_arg_names[0],
                     origin_matmul_output_new_name,
                 )
-            if origin_comm_op.type == "c_allreduce_sum":
+            if origin_comm_op.type == "c_allreduce_sum" or (
+                origin_comm_op.type == "all_reduce"
+                and origin_comm_op.attr("reduce_type")
+                == paddle.distributed.ReduceOp.SUM
+            ):
                 new_comm_var_name = origin_comm_op.input_arg_names[0]
             else:
                 new_comm_var_name = unique_name.generate(
@@ -665,12 +670,12 @@ def _transform_backward(
                         global_block._remove_op(segment[0] + 5)  # scale
                         global_block._remove_op(
                             segment[0] + 4
-                        )  # c_allreduce_sum
+                        )  # all_reduce_sum
                     else:
                         global_block._remove_op(segment[0] + 3)  # scale
                         global_block._remove_op(
                             segment[0] + 2
-                        )  # c_allreduce_sum
+                        )  # all_reduce_sum
                 global_block._sync_with_cpp()
         else:  # not is_first_rank_in tp or sp
             # need to delete the grad op associated with the deleted bias var
@@ -688,7 +693,7 @@ def _transform_backward(
                         global_block._remove_op(segment[0] + 2)  # scale op
                         global_block._remove_op(
                             segment[0] + 1
-                        )  # c_allreduce_sum op
+                        )  # all_reduce_sum op
                     global_block._remove_op(segment[0])
                 global_block._sync_with_cpp()
             else:
@@ -712,11 +717,11 @@ def _transform_backward(
                         )  # scale op for dp
                         global_block._remove_op(
                             segment[0] + 3
-                        )  # c_allreduce_sum op for dp
+                        )  # all_reduce_sum op for dp
                     global_block._remove_op(segment[0] + 2)  # scale op for sp
                     global_block._remove_op(
                         segment[0] + 1
-                    )  # c_allreduce_sum op for sp
+                    )  # all_reduce_sum op for sp
                     global_block._remove_op(
                         segment[0]
                     )  # elementwise_add_grad op
diff --git a/test/deprecated/auto_parallel/test_dist_matmul_deprecated.py b/test/deprecated/auto_parallel/test_dist_matmul_deprecated.py
@@ -169,7 +169,7 @@ def check_row_program(self, main_program, dist_ctx):
         # [0, -1, 1] * [1, -1] --> [0, -1, -1]
         ref_ops = [
             "matmul_v2",
-            "c_allreduce_sum",
+            "all_reduce",
             "reduce_mean",
             "fill_constant",
             "reduce_mean_grad",
@@ -287,7 +287,7 @@ def check_row_program(self, main_program, dist_ctx):
         # [0, -1, 1] * [1, -1] --> [0, -1, -1]
         ref_ops = [
             "matmul_v2",
-            "c_allreduce_sum",
+            "all_reduce",
             "reduce_mean",
             "fill_constant",
             "reduce_mean_grad",
@@ -406,7 +406,7 @@ def check_program(self, main_program, dist_ctx):
             "matmul_v2_grad",
             "c_allreduce_sum",
             "scale",
-            "c_allreduce_sum",
+            "all_reduce",
             "assign",
         ]
         ops = []
diff --git a/test/deprecated/legacy_test/test_auto_parallel_partitioner_deprecated.py b/test/deprecated/legacy_test/test_auto_parallel_partitioner_deprecated.py
@@ -452,7 +452,7 @@ def test_mlp_mp(self):
             'elementwise_add',
             'gelu',
             'matmul_v2',
-            'c_allreduce_sum',
+            'all_reduce',
             'elementwise_add',
             'dropout',
         ]
@@ -545,7 +545,7 @@ def test_mlp_dp_mp(self):
             'elementwise_add',
             'gelu',
             'matmul_v2',
-            'c_allreduce_sum',
+            'all_reduce',
             'elementwise_add',
             'dropout',
         ]
@@ -844,7 +844,7 @@ def test_attn_mp(self):
             'transpose2',
             'reshape2',
             'matmul_v2',
-            'c_allreduce_sum',
+            'all_reduce',
             'elementwise_add',
         ]
         self.assertTrue(dist_ops == ref_ops)
@@ -951,7 +951,7 @@ def test_attn_dp_mp(self):
             'transpose2',
             'reshape2',
             'matmul_v2',
-            'c_allreduce_sum',
+            'all_reduce',
             'elementwise_add',
         ]
         self.assertTrue(dist_ops == ref_ops)
@@ -1320,7 +1320,7 @@ def test_decoder_dp_mp(self):
             'transpose2',
             'reshape2',
             'matmul_v2',
-            'c_allreduce_sum',
+            'all_reduce',
             'elementwise_add',
             'dropout',
             'elementwise_add',
@@ -1329,7 +1329,7 @@ def test_decoder_dp_mp(self):
             'elementwise_add',
             'gelu',
             'matmul_v2',
-            'c_allreduce_sum',
+            'all_reduce',
             'elementwise_add',
             'dropout',
             'elementwise_add',
