From 6cba42e9420d62b8df996ac64fdb891b40790068 Mon Sep 17 00:00:00 2001 From: Wei-Sheng Chin Date: Fri, 9 Oct 2020 15:34:46 -0700 Subject: [PATCH] Avoid inserting other CUDA calls in-between NCCL Send's and Recv's (#5430) * Avoid inserting other CUDA calls in-between NCCL Send's and Recv's * Add a comment * Place CUDA EP on the right device * Fix a warning * Address a comment --- orttraining/orttraining/models/mnist/main.cc | 4 ++- .../cuda/communication/nccl_service.cc | 16 +++++++++-- .../cuda/communication/nccl_service.h | 1 + .../training_ops/cuda/communication/recv.cc | 28 +++++++++++-------- .../training_ops/cuda/communication/send.cc | 12 +++++--- 5 files changed, 42 insertions(+), 19 deletions(-) diff --git a/orttraining/orttraining/models/mnist/main.cc b/orttraining/orttraining/models/mnist/main.cc index 11f0f330b5..52e02df82d 100644 --- a/orttraining/orttraining/models/mnist/main.cc +++ b/orttraining/orttraining/models/mnist/main.cc @@ -79,7 +79,9 @@ Status ParseArguments(int argc, char* argv[], TrainingRunner::Parameters& params #ifdef USE_CUDA bool use_cuda = flags.count("use_cuda") > 0; if (use_cuda) { - params.providers.emplace(kCudaExecutionProvider, CreateExecutionProviderFactory_CUDA(0)); + // Use local rank as device ID of the associated CUDA EP. + OrtDevice::DeviceId device_id = static_cast(MPIContext::GetInstance().GetLocalRank()); + params.providers.emplace(kCudaExecutionProvider, CreateExecutionProviderFactory_CUDA(device_id)); } #endif } catch (const exception& e) { diff --git a/orttraining/orttraining/training_ops/cuda/communication/nccl_service.cc b/orttraining/orttraining/training_ops/cuda/communication/nccl_service.cc index 8c4554db80..c7c3bae818 100644 --- a/orttraining/orttraining/training_ops/cuda/communication/nccl_service.cc +++ b/orttraining/orttraining/training_ops/cuda/communication/nccl_service.cc @@ -233,6 +233,9 @@ void NcclService::Initialize() { // Set device this NCCL communicator runs on. CUDA_CALL(cudaSetDevice(mpi_rank)); + // Create communication stream. + CUDA_CALL(cudaStreamCreateWithFlags(&stream_, cudaStreamNonBlocking)); + // Get NCCL unique ID at rank 0 and broadcast it to all others. ncclUniqueId id; if (mpi_rank == 0) NCCL_CALL(ncclGetUniqueId(&id)); @@ -276,11 +279,11 @@ void NcclService::Launch() { switch (task.type) { case NcclTask::Type::SEND: ORT_ENFORCE(task.peers.size() == 1, "Send can only send data to one rank."); - NCCL_CALL(ncclSend(task.ptr, task.size, ncclChar, task.peers.front(), comm_, nullptr)); + NCCL_CALL(ncclSend(task.ptr, task.size, ncclChar, task.peers.front(), comm_, stream_)); break; case NcclTask::Type::RECV: ORT_ENFORCE(task.peers.size() == 1, "Recv can only send data to one rank."); - NCCL_CALL(ncclRecv(task.ptr, task.size, ncclChar, task.peers.front(), comm_, nullptr)); + NCCL_CALL(ncclRecv(task.ptr, task.size, ncclChar, task.peers.front(), comm_, stream_)); break; default: ORT_NOT_IMPLEMENTED("NCCL service currently only support ncclSend and ncclRecv."); @@ -289,6 +292,13 @@ void NcclService::Launch() { } NCCL_CALL(ncclGroupEnd()); + // Make sure all NCCL computation are done. + // Since the Submit*andWait are blocked by the following "cv_.notify_all()", + // all NCCL Send and Recv called above are all finished before Submit*andWait returning. + // Thus, CUDA operations after Send and Recv won't be inserted by other threads + // when we call NCCL Send's and Recv's. + CUDA_CALL(cudaStreamSynchronize(stream_)); + // This round of communication is done. // We can start waiting for the tasks to be fully scheduled. ++time_; @@ -341,6 +351,8 @@ void NcclService::Terminate() { cv_.wait(lock, [this] { return schedule_.empty() || total_time_ > 0 && time_ == 0; }); } + CUDA_CALL(cudaStreamDestroy(stream_)); + is_running_ = false; worker_.join(); NCCL_CALL(ncclCommDestroy(comm_)); diff --git a/orttraining/orttraining/training_ops/cuda/communication/nccl_service.h b/orttraining/orttraining/training_ops/cuda/communication/nccl_service.h index 0c219d53c8..0186ce2986 100644 --- a/orttraining/orttraining/training_ops/cuda/communication/nccl_service.h +++ b/orttraining/orttraining/training_ops/cuda/communication/nccl_service.h @@ -156,6 +156,7 @@ class NcclService final { std::condition_variable cv_; // Stream for running NCCL. + cudaStream_t stream_; ncclComm_t comm_; // Indicates if NCCL service launched. diff --git a/orttraining/orttraining/training_ops/cuda/communication/recv.cc b/orttraining/orttraining/training_ops/cuda/communication/recv.cc index 0758694216..888e165a50 100644 --- a/orttraining/orttraining/training_ops/cuda/communication/recv.cc +++ b/orttraining/orttraining/training_ops/cuda/communication/recv.cc @@ -67,7 +67,8 @@ void Recv::ReceiveData( profile::NvtxRangeCreator recvRange( "Batch-" + tag + - " Recv-" + std::to_string(src), profile::Color::Green); + " Recv-" + std::to_string(src), + profile::Color::Green); // Begin of major communication tasks. // The first MPI_Recv is not included because we don't want to // count waiting time before setting up the actual communication. @@ -85,15 +86,15 @@ void Recv::ReceiveData( src, static_cast(tag_)}; - // The following NCCL call is equivalent to the following MPI call. User can - // uncomment the MPI call to debug. +// The following NCCL call is equivalent to the following MPI call. User can +// uncomment the MPI call to debug. #if defined(USE_NCCL) && defined(USE_NCCL_P2P) auto& nccl_service = cuda::NcclService::GetInstance(); nccl_service.SubmitRecvAndWait(info_data.buffer, info_data.size, info_data.rank); #else MPI_CHECK(MPI_Recv( - info_data.buffer, info_data.size, MPI_CHAR, - info_data.rank, info_data.tag, MPI_COMM_WORLD, MPI_STATUS_IGNORE)); + info_data.buffer, info_data.size, MPI_CHAR, + info_data.rank, info_data.tag, MPI_COMM_WORLD, MPI_STATUS_IGNORE)); #endif #ifdef ENABLE_NVTX_PROFILE @@ -104,7 +105,8 @@ void Recv::ReceiveData( #ifdef ENABLE_NVTX_PROFILE profile::NvtxRangeCreator memcpyRange( "Batch-" + tag + - " RecvMemcpy-" + std::to_string(src), profile::Color::Green); + " RecvMemcpy-" + std::to_string(src), + profile::Color::Green); // Begin of host-to-device memory copy. memcpyRange.Begin(); #endif @@ -119,11 +121,11 @@ void Recv::ReceiveData( // Copy data out from buffer. #if defined(USE_NCCL) && defined(USE_NCCL_P2P) - CUDA_CALL(cudaMemcpyAsync(tensor->MutableDataRaw(), buffer.get() + tensor_offset_in_bytes, - tensor->SizeInBytes(), cudaMemcpyHostToDevice)); + CUDA_CALL(cudaMemcpy(tensor->MutableDataRaw(), buffer.get() + tensor_offset_in_bytes, + tensor->SizeInBytes(), cudaMemcpyDeviceToDevice)); #else - CUDA_CALL(cudaMemcpyAsync(tensor->MutableDataRaw(), buffer.get() + tensor_offset_in_bytes, - tensor->SizeInBytes(), cudaMemcpyDeviceToDevice)); + CUDA_CALL(cudaMemcpy(tensor->MutableDataRaw(), buffer.get() + tensor_offset_in_bytes, + tensor->SizeInBytes(), cudaMemcpyHostToDevice)); #endif tensor_offset_in_bytes += tensor->SizeInBytes(); } @@ -170,7 +172,8 @@ Status Recv::ComputeInternal(OpKernelContext* ctx) const { profile::NvtxRangeCreator preRange( "Batch-" + tag + - " PreRecv-" + std::to_string(src), profile::Color::Green); + " PreRecv-" + std::to_string(src), + profile::Color::Green); // Begin of preparation for receiving data. preRange.Begin(); #endif @@ -269,7 +272,8 @@ Status Recv::ComputeInternal(OpKernelContext* ctx) const { #ifdef ENABLE_NVTX_PROFILE profile::NvtxRangeCreator postRange( "Batch-" + tag + - " PostRecv-" + std::to_string(src), profile::Color::Green); + " PostRecv-" + std::to_string(src), + profile::Color::Green); postRange.Begin(); #endif diff --git a/orttraining/orttraining/training_ops/cuda/communication/send.cc b/orttraining/orttraining/training_ops/cuda/communication/send.cc index 9b9f3b3fe7..3f54588eb1 100644 --- a/orttraining/orttraining/training_ops/cuda/communication/send.cc +++ b/orttraining/orttraining/training_ops/cuda/communication/send.cc @@ -93,7 +93,8 @@ void Send::SendData( profile::NvtxRangeCreator memcpyRange( "Batch-" + tag + - " SendMemcpy-" + std::to_string(dst), profile::Color::Red); + " SendMemcpy-" + std::to_string(dst), + profile::Color::Red); // Begin of major communication tasks. // The previous MPI_Send's are not included because we don't want to // count waiting time before setting up the actual communication. @@ -125,7 +126,8 @@ void Send::SendData( #ifdef ENABLE_NVTX_PROFILE profile::NvtxRangeCreator sendRange( "Batch-" + tag + - " Send-" + std::to_string(dst), profile::Color::Red); + " Send-" + std::to_string(dst), + profile::Color::Red); // Begin of major communication tasks. // The previous MPI_Send's are not included because we don't want to // count waiting time before setting up the actual communication. @@ -174,7 +176,8 @@ Status Send::ComputeInternal(OpKernelContext* ctx) const { profile::NvtxRangeCreator preRange( "Batch-" + tag + - " PreSend-" + std::to_string(dst), profile::Color::Red); + " PreSend-" + std::to_string(dst), + profile::Color::Red); // Begin of preparation for sending data. This time range includes // the time for sending a scalar. preRange.Begin(); @@ -234,7 +237,8 @@ Status Send::ComputeInternal(OpKernelContext* ctx) const { #ifdef ENABLE_NVTX_PROFILE profile::NvtxRangeCreator postRange( "Batch-" + tag + - " PostSend-" + std::to_string(dst), profile::Color::Red); + " PostSend-" + std::to_string(dst), + profile::Color::Red); // Begin of post communication tasks. postRange.Begin(); #endif