saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-22 19:23:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Avoid inserting other CUDA calls in-between NCCL Send's and Recv's (#5430)

Browse source 
* 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
This commit is contained in:
Wei-Sheng Chin 2020-10-09 15:34:46 -07:00 committed by GitHub
parent dbe7e6623b
commit 6cba42e942
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 42 additions and 19 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
orttraining/orttraining/models/mnist/main.cc
View file

@ -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<OrtDevice::DeviceId>(MPIContext::GetInstance().GetLocalRank());
params.providers.emplace(kCudaExecutionProvider, CreateExecutionProviderFactory_CUDA(device_id));
}
#endif
} catch (const exception& e) {

16
orttraining/orttraining/training_ops/cuda/communication/nccl_service.cc
View file

@ -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_));

1
orttraining/orttraining/training_ops/cuda/communication/nccl_service.h
View file

@ -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.

28
orttraining/orttraining/training_ops/cuda/communication/recv.cc
View file

@ -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<int>(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

12
orttraining/orttraining/training_ops/cuda/communication/send.cc
View file

@ -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