PaddlePaddle · liym27 · Aug 1, 2024 · Jul 24, 2024 · Jul 24, 2024 · Jul 24, 2024
diff --git a/...le/fluid/pir/dialect/distributed/transforms/fuse_allreduce_split_to_reducescatter_pass.cc b/...le/fluid/pir/dialect/distributed/transforms/fuse_allreduce_split_to_reducescatter_pass.cc
@@ -23,7 +23,7 @@
 
 namespace {
 // matmul+c_allreduce_sum_+assign+full+split_with_num+builtin_slice ->
-// c_reducescatter
+// reduce_scatter
 class FusedAllReduceSplitPattern1 : public paddle::drr::DrrPatternBase {
  public:
   std::string name() const override { return "FusedAllReduceSplitPattern1"; }
@@ -76,17 +76,15 @@ class FusedAllReduceSplitPattern1 : public paddle::drr::DrrPatternBase {
 
     paddle::drr::ResultPattern res = pat.ResultPattern();
 
-    const auto &c_reducescatter =
-        res.Op(paddle::dialect::CReducescatterOp::name(),
-               {{"ring_id", pat.Attr("ring_id")},
-                {"nranks", pat.Attr("num")},
-                {"use_calc_stream", pat.Attr("use_calc_stream")}},
+    const auto &reduce_scatter =
+        res.Op(paddle::dialect::ReduceScatterOp::name(),
+               {{"ring_id", pat.Attr("ring_id")}, {"nranks", pat.Attr("num")}},
                {{"execution_stream", pat.Attr("execution_stream")},
                 {"force_record_event", pat.Attr("force_record_event")},
                 {"event_to_record", pat.Attr("event_to_record")},
                 {"events_to_wait", pat.Attr("events_to_wait")}});
 
-    c_reducescatter({&res.Tensor("input_grad_partial")}, {&res.Tensor("out")});
+    reduce_scatter({&res.Tensor("input_grad_partial")}, {&res.Tensor("out")});
   }
 };
 
@@ -128,7 +126,7 @@ class FusedAllReduceSplitPattern1 : public paddle::drr::DrrPatternBase {
 //       |                     |
 //      out1               all_gather
 //       |                     |
-// c_reducescatter        out_g_assign  bias_g
+// reduce_scatter        out_g_assign  bias_g
 //       |                     |
 //      out2    bias       add_grad
 //       |-------|-------------|
@@ -197,11 +195,9 @@ class FusedAllReduceSplitPattern2 : public paddle::drr::DrrPatternBase {
     const auto &res_matmul = res.Op(paddle::dialect::MatmulOp::name(),
                                     {{"transpose_x", pat.Attr("trans_x")},
                                      {"transpose_y", pat.Attr("trans_y")}});
-    const auto &res_c_reducescatter =
-        res.Op(paddle::dialect::CReducescatterOp::name(),
-               {{"ring_id", pat.Attr("ring_id")},
-                {"nranks", pat.Attr("num")},
-                {"use_calc_stream", pat.Attr("use_calc_stream")}},
+    const auto &res_reduce_scatter =
+        res.Op(paddle::dialect::ReduceScatterOp::name(),
+               {{"ring_id", pat.Attr("ring_id")}, {"nranks", pat.Attr("num")}},
                {{"force_record_event", pat.Attr("force_record_event")},
                 {"event_to_record", pat.Attr("event_to_record")},
                 {"events_to_wait", pat.Attr("events_to_wait")}});
@@ -219,7 +215,7 @@ class FusedAllReduceSplitPattern2 : public paddle::drr::DrrPatternBase {
                 {"transpose_y", pat.Attr("mm_g_trans_y")}});
 
     res.Tensor("out1") = res_matmul(res.Tensor("input"), res.Tensor("weight"));
-    res.Tensor("out2") = res_c_reducescatter(res.Tensor("out1"));
+    res.Tensor("out2") = res_reduce_scatter(res.Tensor("out1"));
     res.Tensor("out6") = res_add(res.Tensor("out2"), res.Tensor("bias"));
 
     res_add_grad(

diff --git a/paddle/phi/ops/yaml/legacy/static_ops.yaml b/paddle/phi/ops/yaml/legacy/static_ops.yaml
@@ -720,16 +720,6 @@
     func : reduce
     param: [x, root_id, reduce_type]
 
-- op : reduce_scatter
-  args : (Tensor x, int ring_id = 0, int nranks = 1)
-  output : Tensor(out)
-  infer_meta :
-    func : ReduceScatterInferMeta
-    param: [x, nranks]
-  kernel :
-    func : reduce_scatter
-    param: [x, nranks]
-
 - op : remainder
   args : (Tensor x, Tensor y, int axis = -1)
   output : Tensor (out)

diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -3698,6 +3698,16 @@
   backward : reduce_as_grad
   interfaces : paddle::dialect::InferSymbolicShapeInterface
 
+- op : reduce_scatter
+  args : (Tensor x, int ring_id = 0, int nranks = 1)
+  output : Tensor(out)
+  infer_meta :
+    func : ReduceScatterInferMeta
+    param: [x, nranks]
+  kernel :
+    func : reduce_scatter
+    param: [x, nranks]
+
 - op : reindex_graph
   args : (Tensor x, Tensor neighbors, Tensor count, Tensor hashtable_value, Tensor hashtable_index)
   output : Tensor(reindex_src), Tensor(reindex_dst), Tensor(out_nodes)

diff --git a/python/paddle/distributed/auto_parallel/static/reshard_funcs/p_to_s_reshard_func.py b/python/paddle/distributed/auto_parallel/static/reshard_funcs/p_to_s_reshard_func.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import paddle
+from paddle.distributed.passes.pass_utils import AutoParallelStreamType
 
 from ..process_group import new_process_group
 from .base_reshard_func import (
@@ -76,8 +77,11 @@ def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
 
         num_of_process = len(src_mesh.process_ids)
         group = new_process_group(sorted(src_mesh.process_ids))
-        dst_value = paddle._C_ops.c_reducescatter(
-            src_value, group.id, num_of_process, True
+        dst_value = paddle._C_ops.reduce_scatter(
+            src_value, group.id, num_of_process
+        )
+        dst_value.get_defining_op().set_execution_stream(
+            AutoParallelStreamType.CALC_STREAM.value
         )
 
         # set dist type and dist attr

diff --git a/python/paddle/distributed/passes/auto_parallel_fused_linear_promotion.py b/python/paddle/distributed/passes/auto_parallel_fused_linear_promotion.py
@@ -56,7 +56,7 @@
         "backward": ["elementwise_add_grad", "matmul_v2_grad"],
     },
     {  # MP + SP
-        "forward": ["matmul_v2", "c_reducescatter", "elementwise_add"],
+        "forward": ["matmul_v2", "reduce_scatter", "elementwise_add"],
         "backward": [
             "elementwise_add_grad",
             "c_allreduce_sum",
@@ -76,7 +76,7 @@
         ],
     },
     {  # DP + MP + SP
-        "forward": ["matmul_v2", "c_reducescatter", "elementwise_add"],
+        "forward": ["matmul_v2", "reduce_scatter", "elementwise_add"],
         "backward": [
             "elementwise_add_grad",
             "c_allreduce_sum",
@@ -94,7 +94,7 @@
         "backward": ["elementwise_add_grad", "matmul_v2_grad"],
     },
     {
-        "forward": ["matmul_v2", "c_reducescatter", "cast", "elementwise_add"],
+        "forward": ["matmul_v2", "reduce_scatter", "cast", "elementwise_add"],
         "backward": [
             "elementwise_add_grad",
             "c_allreduce_sum",
@@ -114,7 +114,7 @@
         ],
     },
     {
-        "forward": ["matmul_v2", "c_reducescatter", "cast", "elementwise_add"],
+        "forward": ["matmul_v2", "reduce_scatter", "cast", "elementwise_add"],
         "backward": [
             "elementwise_add_grad",
             "c_allreduce_sum",

diff --git a/test/auto_parallel/test_static_sequence_parallel_pass.py b/test/auto_parallel/test_static_sequence_parallel_pass.py
@@ -191,7 +191,7 @@ def test_decoder_dp_sp(self):
                 allgather_count += 1
 
             # check sequence parallel reducescatter
-            elif op.type == "c_reducescatter":
+            elif op.type == "reduce_scatter":
                 assert (
                     int(op.attr("nranks")) == 4
                 ), "sequence parallel reducescatter error with nranks [{}]".format(

diff --git a/test/collective/test_collective_reduce_scatter_api.py b/test/collective/test_collective_reduce_scatter_api.py
@@ -62,6 +62,29 @@ def test_reduce_scatter_nccl_with_new_comm(self):
                 need_envs={"FLAGS_dynamic_static_unified_comm": "true"},
             )
 
+    def test_allgather_nccl_with_new_comm_pir(self):
+        dtypes_to_test = [
+            "float16",
+            "float32",
+            "float64",
+            "int32",
+            "int64",
+            "int8",
+            "uint8",
+            "bool",
+        ]
+        for dtype in dtypes_to_test:
+            self.check_with_place(
+                "collective_reduce_scatter_api.py",
+                "reduce_sactter",
+                "nccl",
+                dtype=dtype,
+                need_envs={
+                    "FLAGS_dynamic_static_unified_comm": "true",
+                    "FLAGS_enable_pir_in_executor": "1",
+                },
+            )
+
     def test_reduce_scatter_nccl_dygraph(self):
         dtypes_to_test = [
             "float16",

diff --git a/test/distributed_passes/test_fuse_allreduce_split_to_reducescatter_pass.py b/test/distributed_passes/test_fuse_allreduce_split_to_reducescatter_pass.py
@@ -40,7 +40,7 @@ def test_if_with_single_output(self):
         pm.run(ir_program)
         self.assertEqual(ir_program.global_block().num_ops(), 6)
         self.assertEqual(
-            ir_program.global_block().ops[-2].name(), "pd_op.c_reducescatter"
+            ir_program.global_block().ops[-2].name(), "pd_op.reduce_scatter"
         )