After nccl add event

7 years ago · 9b1f4d5d62
parent feb569f8ea
commit 9b1f4d5d62
1 changed files with 5 additions and 2 deletions
--- a/paddle/fluid/framework/parallel_executor.cc
+++ b/paddle/fluid/framework/parallel_executor.cc
@ -402,10 +402,13 @@ struct NCCLAllReduceOpHandle : public OpHandle {
        platform::dynload::ncclAllReduce(
            buffer, buffer, numel, static_cast<ncclDataType_t>(dtype), ncclSum,
            nccl_ctx.comm, nccl_ctx.stream());
-        PADDLE_ENFORCE(cudaEventRecord(events_[dev_id], nccl_ctx.stream()));
      }
-
      platform::dynload::ncclGroupEnd();
+
+      for (auto &ev : events_) {
+        PADDLE_ENFORCE(cudaEventRecord(
+            ev.second, member_->communication_streams_.at(ev.first).stream()));
+      }
    }
  }