[fluid_ops] fleet modify c_allreduce_sum (#72092)

co63oc · web-flow · commit 430a255f2e48 · 2025-04-08T10:55:07.000+08:00
* Fix

* Fix

* Fix
diff --git a/python/paddle/distributed/fleet/meta_optimizers/common.py b/python/paddle/distributed/fleet/meta_optimizers/common.py
@@ -117,12 +117,12 @@ def _add_sync_by_allreduce(block):
                 },
             )
             block.append_op(
-                type='c_allreduce_sum',
-                inputs={'X': [sync_var]},
-                outputs={'Out': [sync_var]},
+                type='all_reduce',
+                inputs={'x': [sync_var]},
+                outputs={'out': [sync_var]},
                 attrs={
                     'ring_id': global_ring_id,
-                    'use_calc_stream': True,
+                    'reduce_type': paddle.distributed.ReduceOp.SUM,
                     OP_ROLE_KEY: OpRole.Forward,
                 },
             )
diff --git a/python/paddle/distributed/fleet/meta_optimizers/localsgd_optimizer.py b/python/paddle/distributed/fleet/meta_optimizers/localsgd_optimizer.py
@@ -156,11 +156,12 @@ def communicate():
                     )
                     ring_id = (ring_id + 1) % self.nrings
                     sub_block.append_op(
-                        type='c_allreduce_sum',
-                        inputs={'X': [param]},
-                        outputs={'Out': [param]},
+                        type='all_reduce',
+                        inputs={'x': [param]},
+                        outputs={'out': [param]},
                         attrs={
                             'ring_id': ring_id,
+                            'reduce_type': paddle.distributed.ReduceOp.SUM,
                             OP_ROLE_KEY: OpRole.Optimize,
                         },
                     )
@@ -278,13 +279,13 @@ def init_snapshot_vars(self, startup_program, param2snapshot):
 
     def _generate_avg_loss(self, program_block, loss, avg_loss):
         program_block.append_op(
-            type='c_allreduce_sum',
-            inputs={'X': [loss]},
-            outputs={'Out': [avg_loss]},
+            type='all_reduce',
+            inputs={'x': [loss]},
+            outputs={'out': [avg_loss]},
             attrs={
                 'ring_id': 0,
                 OP_ROLE_KEY: OpRole.Optimize,
-                'use_calc_stream': True,
+                'reduce_type': paddle.distributed.ReduceOp.SUM,
             },
         )
         program_block.append_op(
@@ -407,11 +408,12 @@ def communicate():
                     )
                     ring_id = (ring_id + 1) % self.nrings
                     sub_block.append_op(
-                        type='c_allreduce_sum',
-                        inputs={'X': [param]},
-                        outputs={'Out': [param]},
+                        type='all_reduce',
+                        inputs={'x': [param]},
+                        outputs={'out': [param]},
                         attrs={
                             'ring_id': ring_id,
+                            'reduce_type': paddle.distributed.ReduceOp.SUM,
                             OP_ROLE_KEY: OpRole.Optimize,
                         },
                     )
diff --git a/python/paddle/distributed/fleet/meta_optimizers/pipeline_optimizer.py b/python/paddle/distributed/fleet/meta_optimizers/pipeline_optimizer.py
@@ -308,12 +308,12 @@ def _insert_allreduce_ops(self, ring_id):
 
                     block._insert_op(
                         first_optimize_op_idx + offset,
-                        type='c_allreduce_sum',
-                        inputs={'X': grad},
-                        outputs={'Out': grad},
+                        type='all_reduce',
+                        inputs={'x': grad},
+                        outputs={'out': grad},
                         attrs={
                             'ring_id': ring_id,
-                            'use_calc_stream': True,
+                            'reduce_type': paddle.distributed.ReduceOp.SUM,
                             OP_ROLE_KEY: OpRole.Optimize,
                         },
                     )
diff --git a/python/paddle/distributed/fleet/meta_optimizers/raw_program_optimizer.py b/python/paddle/distributed/fleet/meta_optimizers/raw_program_optimizer.py
@@ -13,6 +13,7 @@
 
 import os
 
+import paddle
 from paddle import static
 from paddle.base import core
 from paddle.framework.ir import apply_build_strategy
@@ -257,11 +258,12 @@ def _insert_allreduce_ops_for_gm(self, gm_block):
         for i, (p, g) in enumerate(zip(param_vars, grad_vars)):
             gm_block._insert_op(
                 first_optimize_op_idx + insert_op_num,
-                type="c_allreduce_sum",
-                inputs={'X': g},
-                outputs={'Out': g},
+                type="all_reduce",
+                inputs={'x': g},
+                outputs={'out': g},
                 attrs={
                     'ring_id': ring_id,
+                    'reduce_type': paddle.distributed.ReduceOp.SUM,
                     OP_ROLE_KEY: OpRole.Backward,
                 },
             )
@@ -333,11 +335,12 @@ def _insert_allreduce_ops(self):
 
                     block._insert_op(
                         idx + offset,
-                        type='c_allreduce_sum',
-                        inputs={'X': grad},
-                        outputs={'Out': grad},
+                        type='all_reduce',
+                        inputs={'x': grad},
+                        outputs={'out': grad},
                         attrs={
                             'ring_id': ring_id,
+                            'reduce_type': paddle.distributed.ReduceOp.SUM,
                             OP_ROLE_KEY: OpRole.Backward,
                         },
                     )
@@ -440,12 +443,12 @@ def get_after_idx_of_fuse_group(grad_param_segments):
             )
             block._insert_op_without_sync(
                 after_idx + 1,
-                type='c_allreduce_sum',
-                inputs={'X': fused_var},
-                outputs={'Out': fused_var},
+                type='all_reduce',
+                inputs={'x': fused_var},
+                outputs={'out': fused_var},
                 attrs={
                     'ring_id': ring_id,
-                    'use_calc_stream': self.calc_comm_same_stream,
+                    'reduce_type': paddle.distributed.ReduceOp.SUM,
                     OP_ROLE_KEY: OpRole.Backward,
                 },
             )
@@ -462,8 +465,14 @@ def get_after_idx_of_fuse_group(grad_param_segments):
             for i in range(len(grad_param_segments)):
                 while (
                     block.ops[idx].type != 'c_allreduce_sum'
-                    or fused_vars[i].name not in block.ops[idx].input_arg_names
-                ):
+                    and (
+                        not (
+                            block.ops[idx].type == 'all_reduce'
+                            and block.ops[idx].attr('reduce_type')
+                            == paddle.distributed.ReduceOp.SUM
+                        )
+                    )
+                ) or fused_vars[i].name not in block.ops[idx].input_arg_names:
                     idx += 1
                 grad_segment, param_segment = grad_param_segments[i]
                 for grad in grad_segment:
diff --git a/python/paddle/distributed/fleet/meta_optimizers/sharding/gradient_clip_helper.py b/python/paddle/distributed/fleet/meta_optimizers/sharding/gradient_clip_helper.py
@@ -248,13 +248,13 @@ def _insert_allreduce(block, ring_ids, idx, var):
             idx = idx + 1
             block._insert_op_without_sync(
                 idx,
-                type='c_allreduce_sum',
-                inputs={'X': var},
-                outputs={'Out': var},
+                type='all_reduce',
+                inputs={'x': var},
+                outputs={'out': var},
                 attrs={
                     'ring_id': ring_id,
                     'op_namescope': "/gradient_clip_model_parallelism",
-                    'use_calc_stream': True,
+                    'reduce_type': paddle.distributed.ReduceOp.SUM,
                     OP_ROLE_KEY: OpRole.Optimize,
                 },
             )
diff --git a/python/paddle/distributed/fleet/meta_optimizers/sharding/utils.py b/python/paddle/distributed/fleet/meta_optimizers/sharding/utils.py
@@ -167,9 +167,16 @@ def check_allreduce_sum(block, shard, sharding_ring_id, dp_ring_id=-1):
                 ):
                     dp_grads_status[var_name] = 1
         # check sharding allreduce and  reduce but skip megatron allreduce
-        elif op.type == "c_allreduce_sum" or (
-            op.type == "reduce"
-            and op.desc.attr("reduce_type") == dist.ReduceOp.SUM
+        elif (
+            op.type == "c_allreduce_sum"
+            or (
+                op.type == "all_reduce"
+                and op.desc.attr("reduce_type") == dist.ReduceOp.SUM
+            )
+            or (
+                op.type == "reduce"
+                and op.desc.attr("reduce_type") == dist.ReduceOp.SUM
+            )
         ):
             if not op.all_attrs()["use_calc_stream"]:
                 var_name = op.desc.input_arg_names()[0]
@@ -517,12 +524,12 @@ def insert_fused_allreduce_ops(
     for fused_var in fused_vars:
         block._insert_op_without_sync(
             insert_idx + insert_num,
-            type='c_allreduce_sum',
-            inputs={'X': fused_var},
-            outputs={'Out': fused_var},
+            type='all_reduce',
+            inputs={'x': fused_var},
+            outputs={'out': fused_var},
             attrs={
                 'ring_id': ring_id,
-                'use_calc_stream': use_calc_stream,
+                'reduce_type': paddle.distributed.ReduceOp.SUM,
                 OP_ROLE_KEY: op_role,
             },
         )
diff --git a/python/paddle/distributed/fleet/meta_optimizers/sharding_optimizer.py b/python/paddle/distributed/fleet/meta_optimizers/sharding_optimizer.py
@@ -14,6 +14,7 @@
 
 import os
 
+import paddle
 from paddle.base import core
 from paddle.incubate.optimizer import PipelineOptimizer
 from paddle.static import (
@@ -998,6 +999,9 @@ def _prune_main_program(self, block, shard, rings):
             if (
                 op.type == "c_allreduce_sum"
                 and op.attr('use_model_parallel') is False
+            ) or (
+                op.type == "all_reduce"
+                and op.attr('reduce_type') == paddle.distributed.ReduceOp.SUM
             ):
                 assert len(output_names) == 1
                 output_name = output_names[0]
@@ -1018,6 +1022,7 @@ def _prune_main_program(self, block, shard, rings):
         for idx, op in reversed(list(enumerate(block.ops))):
             if op.type in [
                 "c_allreduce_sum",
+                "all_reduce",
                 "c_sync_comm_stream",
                 "c_calc_comm_stream",
                 "c_gen_nccl_id",
@@ -1881,12 +1886,12 @@ def _true_apply_gradient(self):
             for grad, merged_grad in self._grad2merged_grad.items():
                 merged_grad_var = main_block.var(merged_grad)
                 cur_block.append_op(
-                    type='c_allreduce_sum',
+                    type='all_reduce',
                     inputs={'X': merged_grad_var},
                     outputs={'Out': merged_grad_var},
                     attrs={
                         'ring_id': self.dp_ring_id,
-                        'use_calc_stream': True,
+                        'reduce_type': paddle.distributed.ReduceOp.SUM,
                         OP_ROLE_KEY: OpRole.Optimize,
                     },
                 )
diff --git a/python/paddle/distributed/fleet/meta_optimizers/tensor_parallel_optimizer.py b/python/paddle/distributed/fleet/meta_optimizers/tensor_parallel_optimizer.py
@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 
+import paddle
 from paddle import static
 
 from .common import (
@@ -244,11 +245,12 @@ def _insert_allreduce_ops(self, loss, ring_id):
 
                     block._insert_op(
                         offset,
-                        type='c_allreduce_sum',
-                        inputs={'X': grad},
-                        outputs={'Out': grad},
+                        type='all_reduce',
+                        inputs={'x': grad},
+                        outputs={'out': grad},
                         attrs={
                             'ring_id': ring_id,
+                            'reduce_type': paddle.distributed.ReduceOp.SUM,
                             OP_ROLE_KEY: OpRole.Backward,
                         },
                     )
diff --git a/python/paddle/incubate/distributed/fleet/parameter_server/pslib/__init__.py b/python/paddle/incubate/distributed/fleet/parameter_server/pslib/__init__.py
@@ -1263,7 +1263,7 @@ def _remove_collective_op_for_embedding(self, loss, table_name):
         table_name = [name + "@GRAD" for name in table_name]
         need_remove_op_index = []
         block = loss.block.program.global_block()
-        collective_ops = ["c_sync_calc_stream", "c_allreduce_sum"]
+        collective_ops = ["c_sync_calc_stream", "c_allreduce_sum", "all_reduce"]
         for ids, op in list(enumerate(block.ops)):
             if op.type in collective_ops:
                 if op.input("X")[0] in table_name:
diff --git a/python/paddle/static/amp/fp16_lists.py b/python/paddle/static/amp/fp16_lists.py
@@ -263,6 +263,7 @@ def _update_list(self):
     'c_identity',
     'c_concat',
     'c_allreduce_sum',
+    'all_reduce',
     'concat',
     'split',
     'fused_feedforward',
