Fix

Author: co63oc

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

@@ -30,6 +30,7 @@
     "broadcast",
     "all_gather",
     "c_allreduce_sum",
+    "all_reduce",
     "c_identity",
 ]
 NON_COMP_TYPE = ["while", *COMM_OP_TYPE]
@@ -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,7 +423,7 @@ def build_dp_costs(
     if not has_found:
         return
 
-    c_allreduce_sum_descs = build_comm_desc_from_dist_op(
+    all_reduce_sum_descs = build_comm_desc_from_dist_op(
         "c_allreduce_sum",
         dist_op,
         ctx,
@@ -428,7 +435,7 @@ def build_dp_costs(
         _g_op_cost_factory["c_allreduce_sum"],
         ctx,
         processes,
-        c_allreduce_sum_descs,
+        all_reduce_sum_descs,
         cluster,
         is_dp=True,
     )
@@ -787,17 +794,27 @@ def comm_count(self):
                 vars = self.op.block.vars
                 # NOTE: The tensor communicated input_name is "X" in default. Otherwise, this function should be overridden
                 try:
-                    var_name = self.op.input("X")[0]
+                    if self.op.type != "all_reduce":
+                        var_name = self.op.input("X")[0]
+                    else:
+                        var_name = self.op.input("x")[0]
                 except:
-                    var_name = self.op.output("Out")[0]
+                    if self.op.type != "all_reduce":
+                        var_name = self.op.output("Out")[0]
+                    else:
+                        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]
+                if "op" in self.op_desc and self.op_desc["op"] == "all_reduce":
+                    dtype = self.op_desc["inputs"]["x"][0][0]
+                    shape = self.op_desc["inputs"]["x"][0][1]
+                else:
+                    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:

python/paddle/distributed/passes/auto_parallel_data_parallel_optimization.py

@@ -280,7 +280,6 @@ def _comms_overlap_calc(self):
         # comm wait calc to finish
         for idx, op in reversed(list(enumerate(block.ops))):
             if is_data_parallel_reduce_op(op):
-                assert op.has_attr('use_calc_stream')
                 assert op.has_attr('ring_id')
 
                 op._set_attr('use_calc_stream', False)
@@ -492,6 +491,7 @@ def _update_program(self, grad_groups):
 
             allreduce_op = block.ops[group.allreduce_op_idx]
             assert allreduce_op.type in [
+                'all_reduce',
                 'c_allreduce_avg',
                 'c_allreduce_sum',
             ], f"should found c_allreduce_avg or c_allreduce_sum op but found {allreduce_op}"