【Comm】Fix execution stream for comm op (#67176)

liym27 · web-flow · commit fffb7f771a08 · 2024-08-13T14:16:07.000+08:00
diff --git a/paddle/fluid/framework/new_executor/instruction/instruction_util.cc b/paddle/fluid/framework/new_executor/instruction/instruction_util.cc
@@ -133,26 +133,38 @@ phi::DeviceContext* ParseDeviceContext(pir::Operation* op,
       }
       return dev_ctx;
     }
-    if (FLAGS_dynamic_static_unified_comm) {
-      if (op_attributes.count("ring_id") != 0) {
-        int ring_id =
-            op_attributes.at("ring_id").dyn_cast<pir::Int32Attribute>().data();
-        const auto& comm_context_manager =
-            phi::distributed::CommContextManager::GetInstance();
-        if (comm_context_manager.Has(std::to_string(ring_id))) {
-          auto comm_context = comm_context_manager.Get(std::to_string(ring_id));
-          dev_ctx = static_cast<platform::DeviceContext*>(
-              static_cast<phi::distributed::NCCLCommContext*>(comm_context)
-                  ->GetDevContext());
-          dev_ctx->SetCommContext(comm_context);
-          if (op_name.compare(paddle::dialect::ReduceScatterOp::name()) == 0 ||
-              op_name.compare(paddle::dialect::AllGatherOp::name()) == 0) {
-            return dev_ctx;
+
+    // handle comm op
+    if (op_attributes.count("ring_id") != 0 &&
+        FLAGS_dynamic_static_unified_comm) {
+      int ring_id =
+          op_attributes.at("ring_id").dyn_cast<pir::Int32Attribute>().data();
+      const auto& comm_context_manager =
+          phi::distributed::CommContextManager::GetInstance();
+      if (comm_context_manager.Has(std::to_string(ring_id))) {
+        auto comm_context = comm_context_manager.Get(std::to_string(ring_id));
+        dev_ctx = static_cast<platform::DeviceContext*>(
+            static_cast<phi::distributed::NCCLCommContext*>(comm_context)
+                ->GetDevContext());
+        dev_ctx->SetCommContext(comm_context);
+        if (op_name.compare(paddle::dialect::ReduceScatterOp::name()) == 0 ||
+            op_name.compare(paddle::dialect::AllGatherOp::name()) == 0) {
+          if (phi::is_gpu_place(place) && execution_stream == kDefaultStream) {
+            if (origin_dev_ctx != nullptr) {
+              auto default_stream =
+                  static_cast<phi::GPUContext*>(origin_dev_ctx)->cuda_stream();
+              static_cast<phi::GPUContext*>(dev_ctx)->SetCUDAStream(
+                  default_stream, false);
+            } else {
+              VLOG(3) << "op " << op_name << " ring_id " << ring_id
+                      << " origin_dev_ctx is nullptr";
+            }
           }
-        } else {
-          VLOG(10) << "ring_id " << ring_id
-                   << " not found in comm_context_manager for op " << op_name;
+          return dev_ctx;
         }
+      } else {
+        VLOG(3) << "ring_id " << ring_id
+                << " not found in comm_context_manager for op " << op_name;
       }
     }
 #endif
diff --git a/python/paddle/distributed/auto_parallel/static/operators/dist_pnorm.py b/python/paddle/distributed/auto_parallel/static/operators/dist_pnorm.py
@@ -15,7 +15,7 @@
 import copy
 
 from paddle.common_ops_import import check_dtype, check_variable_and_dtype
-from paddle.distributed.passes.pass_utils import AutoParallelStreamType
+from paddle.distributed.utils.stream_utils import ExecutionStreamType
 from paddle.framework import core
 from paddle.static import Operator
 
@@ -243,7 +243,7 @@ def forward(ctx, *args, **kwargs):
             allgather_out.name, allgather_out_dist_attr.dims_mapping
         )
         allgather_op_dist_attr.execution_stream = (
-            AutoParallelStreamType.CALC_STREAM.value
+            ExecutionStreamType.DefaultStream.value
         )
         ctx.set_op_dist_attr_for_program(all_gather_op, allgather_op_dist_attr)
 
diff --git a/python/paddle/distributed/auto_parallel/static/reshard.py b/python/paddle/distributed/auto_parallel/static/reshard.py
@@ -20,7 +20,7 @@
 
 import paddle
 from paddle.distributed.fleet.meta_optimizers.common import OpRole
-from paddle.distributed.passes.pass_utils import AutoParallelStreamType
+from paddle.distributed.utils.stream_utils import ExecutionStreamType
 from paddle.framework import LayerHelper, OpProtoHolder, Program, core
 from paddle.utils import unique_name
 
@@ -702,7 +702,7 @@ def insert_allgather_op(
         )
         allgather_op._set_attr('op_namescope', "/auto_parallel/reshard")
         allgather_op.dist_attr.execution_stream = (
-            AutoParallelStreamType.CALC_STREAM.value
+            ExecutionStreamType.DefaultStream.value
         )
         idx_offset += 1
 
diff --git a/python/paddle/distributed/communication/stream/all_gather.py b/python/paddle/distributed/communication/stream/all_gather.py
@@ -27,6 +27,8 @@
     from paddle.base.core import task
     from paddle.distributed.communication.group import Group
 
+from paddle.distributed.utils.stream_utils import ExecutionStreamType
+
 
 def _all_gather_into_tensor_in_dygraph(
     out_tensor: Tensor,
@@ -130,7 +132,7 @@ def _all_gather_in_static_mode(
         },
     )
     if sync_op:
-        op.dist_attr.execution_stream = "default"
+        op.dist_attr.execution_stream = ExecutionStreamType.DefaultStream.value
     tensor_list.clear()
     # 0-D use stack/unstack while others use concat/split
     if len(tensor.shape) == 0:
diff --git a/python/paddle/distributed/utils/stream_utils.py b/python/paddle/distributed/utils/stream_utils.py
@@ -0,0 +1,19 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from enum import Enum
+
+
+class ExecutionStreamType(Enum):
+    DefaultStream = "DefaultStream"

Original file line number	Diff line number	Diff line change
`@@ -15,7 +15,7 @@`
`15`	`15`	`import copy`
`16`	`16`
`17`	`17`	`from paddle.common_ops_import import check_dtype, check_variable_and_dtype`
`18`		`-from paddle.distributed.passes.pass_utils import AutoParallelStreamType`
	`18`	`+from paddle.distributed.utils.stream_utils import ExecutionStreamType`
`19`	`19`	`from paddle.framework import core`
`20`	`20`	`from paddle.static import Operator`
`21`	`21`
`@@ -243,7 +243,7 @@ def forward(ctx, args, *kwargs):`
`243`	`243`	`allgather_out.name, allgather_out_dist_attr.dims_mapping`
`244`	`244`	`)`
`245`	`245`	`allgather_op_dist_attr.execution_stream = (`
`246`		`- AutoParallelStreamType.CALC_STREAM.value`
	`246`	`+ ExecutionStreamType.DefaultStream.value`
`247`	`247`	`)`
`248`	`248`	`ctx.set_op_dist_attr_for_program(all_gather_op, allgather_op_dist_attr)`
`249`	`249`
Original file line number	Diff line number	Diff line change
`@@ -20,7 +20,7 @@`
`20`	`20`
`21`	`21`	`import paddle`
`22`	`22`	`from paddle.distributed.fleet.meta_optimizers.common import OpRole`
`23`		`-from paddle.distributed.passes.pass_utils import AutoParallelStreamType`
	`23`	`+from paddle.distributed.utils.stream_utils import ExecutionStreamType`
`24`	`24`	`from paddle.framework import LayerHelper, OpProtoHolder, Program, core`
`25`	`25`	`from paddle.utils import unique_name`
`26`	`26`
`@@ -702,7 +702,7 @@ def insert_allgather_op(`
`702`	`702`	`)`
`703`	`703`	`allgather_op._set_attr('op_namescope', "/auto_parallel/reshard")`
`704`	`704`	`allgather_op.dist_attr.execution_stream = (`
`705`		`- AutoParallelStreamType.CALC_STREAM.value`
	`705`	`+ ExecutionStreamType.DefaultStream.value`
`706`	`706`	`)`
`707`	`707`	`idx_offset += 1`
`708`	`708`