Fix

Author: co63oc

python/paddle/distributed/auto_parallel/static/cost/base_cost.py

@@ -29,6 +29,7 @@
     "recv_v2",
     "broadcast",
     "all_gather",
+    "all_reduce",
     "c_allreduce_sum",
     "c_identity",
 ]
@@ -311,7 +312,10 @@ def build_comm_desc_from_dist_op(
                 input_list.append((var.dtype, shape))
 
             # NOTE: The input_name of comm ops used usually is X.
-            desc["inputs"] = {"X": input_list}
+            if op_type == "all_reduce":
+                desc["inputs"] = {"x": input_list}
+            else:
+                desc["inputs"] = {"X": input_list}
 
             # Get comm group by parallel_axis or the given group_ranks.
             if parallel_axis is not None:
@@ -349,7 +353,10 @@ def build_comm_desc(op_type, group_ranks, dtype, shape, attrs=None):
     desc = {}
     desc["op"] = op_type
     desc["group_ranks"] = group_ranks
-    desc["inputs"] = {"X": [(dtype, shape)]}
+    if op_type == "all_reduce":
+        desc["inputs"] = {"x": [(dtype, shape)]}
+    else:
+        desc["inputs"] = {"X": [(dtype, shape)]}
     desc["attrs"] = attrs
     return desc
 
@@ -416,19 +423,19 @@ def build_dp_costs(
     if not has_found:
         return
 
-    c_allreduce_sum_descs = build_comm_desc_from_dist_op(
-        "c_allreduce_sum",
+    all_reduce_sum_descs = build_comm_desc_from_dist_op(
+        "all_reduce",
         dist_op,
         ctx,
         var_names,
         attrs=attrs,
         parallel_axis=parallel_axis,
     )
     comm_cost_list = build_comm_costs_from_descs(
-        _g_op_cost_factory["c_allreduce_sum"],
+        _g_op_cost_factory["all_reduce"],
         ctx,
         processes,
-        c_allreduce_sum_descs,
+        all_reduce_sum_descs,
         cluster,
         is_dp=True,
     )

python/paddle/distributed/auto_parallel/static/cost/comm_op_cost.py

@@ -23,7 +23,7 @@
 
 @register_op_cost
 class AllreduceSumOpCost(CommOpCost):
-    OP_TYPE = "c_allreduce_sum"
+    OP_TYPE = "all_reduce"
 
     def __init__(self, op=None, op_desc=None, comm_context=None):
         super().__init__(op=op, op_desc=op_desc, comm_context=comm_context)
@@ -82,6 +82,38 @@ def calc_time_tree(self):
 
         return time
 
+    @property
+    def comm_count(self):
+        from ..reshard import get_var_with_recursion
+
+        if self._comm_count is None:
+            dtype = None
+            shape = None
+            if self.op is not None:
+                vars = self.op.block.vars
+                try:
+                    var_name = self.op.input("x")[0]
+                except:
+                    var_name = self.op.output("out")[0]
+                var = get_var_with_recursion(
+                    var_name, self.op.block, self.op.block.program
+                )
+                dtype = var.dtype
+                shape = var.shape
+            elif self.op_desc is not None:
+                dtype = self.op_desc["inputs"]["x"][0][0]
+                shape = self.op_desc["inputs"]["x"][0][1]
+
+            factor = None
+            if dtype == paddle.float32 or dtype == paddle.int32:
+                factor = 4
+            else:
+                raise ValueError(f"Unsupported comm dtype {dtype}")
+            comm_count = int(np.prod(shape)) * factor
+            self._comm_count = comm_count
+
+        return self._comm_count
+
 
 @register_op_cost
 class AllgatherOpCost(CommOpCost):

python/paddle/distributed/auto_parallel/static/mapper.py

@@ -104,7 +104,7 @@ def get_comm_volume(comm_op, src_rank, tgt_rank):
             new_tensor_shape.append(val)
     tensor_size = functools.reduce(operator.mul, new_tensor_shape, 1)
     tensor_bytes = tensor_size * get_dtype_bytes(tensor.dtype)
-    if "c_allreduce" in comm_op_type:
+    if "c_allreduce" in comm_op_type or "all_reduce" in comm_op_type:
         comm_volume = 2 * tensor_bytes
     elif "all_gather" in comm_op_type:
         comm_volume = tensor_bytes

python/paddle/distributed/auto_parallel/static/operators/common.py

@@ -511,17 +511,17 @@ def sync_and_scale_gradients(dist_ctx, op, groups, allreduce_var_names):
     dist_op_context = dist_ctx.dist_op_context
     main_block = dist_op_context.work_block
 
-    allreduce_type = "c_allreduce_sum"
+    op_type = dist.ReduceOp.SUM
     need_scale = dist_ctx.gradient_scale
     scale_using_allreduce_avg = dist_ctx.gradient_scale_using_allreduce_avg
 
-    # With nccl_version > 2.10.00, we can use c_allreduce_avg to replace c_allreduce_sum and eliminate the scale op.
+    # With nccl_version > 2.10.00, we can use all_reduce_avg to replace all_reduce_sum and eliminate the scale op.
     if (
         need_scale
         and scale_using_allreduce_avg
         and int(paddle.version.nccl()) > 21000
     ):
-        allreduce_type = "c_allreduce_avg"
+        op_type = dist.ReduceOp.AVG
         need_scale = False
 
     for group in groups:
@@ -531,12 +531,12 @@ def sync_and_scale_gradients(dist_ctx, op, groups, allreduce_var_names):
             added_ops = []
             grad_var = main_block.var(var_name)
             allreduce_op = main_block.append_op(
-                type=allreduce_type,
-                inputs={'X': [grad_var]},
-                outputs={'Out': [grad_var]},
+                type='all_reduce',
+                inputs={'x': [grad_var]},
+                outputs={'out': [grad_var]},
                 attrs={
                     'ring_id': group.id,
-                    'use_calc_stream': True,
+                    'op_type': op_type,
                     OP_ROLE_KEY: OpRole.Backward,
                 },
             )
@@ -670,9 +670,11 @@ def is_data_parallel_scale_op(op):
 
 
 def is_data_parallel_reduce_op(op):
-    is_allreduce_op = op.type in [
-        "c_allreduce_sum",
-        "c_allreduce_avg",
+    is_allreduce_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,
@@ -695,7 +697,8 @@ def is_amp_flag_sync_op(op):
 
 def is_global_norm_sync_op(op):
     return (
-        op.type == "c_allreduce_sum"
+        op.type == "all_reduce"
+        and op.desc.attr("op_type") == dist.ReduceOp.SUM
         and op.desc.has_attr("op_namescope")
         and SyncMode.GlobalNormSync in op.desc.attr("op_namescope")
     )

python/paddle/distributed/auto_parallel/static/operators/dist_default.py

@@ -13,6 +13,7 @@
 # limitations under the License
 
 
+import paddle.distributed as dist
 from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY, OpRole
 
 from ..completion import contains_spmd_rule, get_phi_spmd_rule
@@ -67,12 +68,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,
+                'op_type': dist.ReduceOp.SUM,
                 OP_ROLE_KEY: OpRole.Backward,
             },
         )

python/paddle/distributed/auto_parallel/static/operators/dist_embedding.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License
 
+import paddle.distributed as dist
 from paddle.common_ops_import import check_variable_and_dtype
 from paddle.distributed.auto_parallel.static.cost.comm_op_cost import (
     AllreduceSumOpCost,
@@ -246,10 +247,10 @@ def calc_fwd_cost(self, dist_op, ctx, cluster):
         parallel_axis = dist_op.dist_attr.get_input_dims_mapping(
             serial_op.input("W")[0]
         )[0]
-        attrs = {"use_calc_stream": True, "use_model_parallel": True}
+        attrs = {"op_type": dist.ReduceOp.SUM}
         var_names = serial_op.output("Out")
-        c_allreduce_sum_desc_mapping = build_comm_desc_from_dist_op(
-            "c_allreduce_sum",
+        all_reduce_sum_desc_mapping = build_comm_desc_from_dist_op(
+            "all_reduce",
             dist_op,
             ctx,
             var_names,
@@ -261,7 +262,7 @@ def calc_fwd_cost(self, dist_op, ctx, cluster):
             AllreduceSumOpCost,
             ctx,
             processes,
-            c_allreduce_sum_desc_mapping,
+            all_reduce_sum_desc_mapping,
             cluster,
         )
 
@@ -510,23 +511,22 @@ def forward(ctx, *args, **kwargs):
         naive_copy_op_dist_attr_for_program(c_embedding_op, src_op, ctx)
 
         # use_model_parallel
-        c_allreduce_sum_op = main_block.append_op(
-            type='c_allreduce_sum',
-            inputs={'X': [Out_var]},
-            outputs={'Out': [Out_var]},
+        all_reduce_sum_op = main_block.append_op(
+            type='all_reduce',
+            inputs={'x': [Out_var]},
+            outputs={'out': [Out_var]},
             attrs={
                 'ring_id': group.id,
-                'use_calc_stream': True,
-                'use_model_parallel': True,
+                'op_type': dist.ReduceOp.SUM,
                 OP_ROLE_KEY: src_op.attr('op_role'),
             },
         )
-        c_allreduce_sum_op._set_attr(
+        all_reduce_sum_op._set_attr(
             'op_namescope', '/' + ParallelMode.TensorParallel
         )
         # allreduce
         set_comm_op_dist_attr_for_program(
-            c_allreduce_sum_op,
+            all_reduce_sum_op,
             op_dist_attr.process_mesh,
             out_var_dist_attr,
             ctx,