After nccl add event

helinwang-patch-1
Yu Yang 7 years ago
parent feb569f8ea
commit 9b1f4d5d62

@ -402,10 +402,13 @@ struct NCCLAllReduceOpHandle : public OpHandle {
platform::dynload::ncclAllReduce( platform::dynload::ncclAllReduce(
buffer, buffer, numel, static_cast<ncclDataType_t>(dtype), ncclSum, buffer, buffer, numel, static_cast<ncclDataType_t>(dtype), ncclSum,
nccl_ctx.comm, nccl_ctx.stream()); nccl_ctx.comm, nccl_ctx.stream());
PADDLE_ENFORCE(cudaEventRecord(events_[dev_id], nccl_ctx.stream()));
} }
platform::dynload::ncclGroupEnd(); platform::dynload::ncclGroupEnd();
for (auto &ev : events_) {
PADDLE_ENFORCE(cudaEventRecord(
ev.second, member_->communication_streams_.at(ev.first).stream()));
}
} }
} }

Loading…
Cancel
Save