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

Support Send and Recv for old NCCL versions (#5097)

Browse source 
If NCCL version < 2.7, MPI is sued. Otherwise, we use NCCL Send and Recv.
This commit is contained in:
Wei-Sheng Chin 2020-09-09 20:58:05 -07:00 committed by GitHub
parent 09a6ce6bc0
commit 9ba56dcfed
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 27 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
cmake/CMakeLists.txt
View file

@ -1172,6 +1172,16 @@ if (onnxruntime_ENABLE_TRAINING)
message(STATUS "NCCL_MINOR_VERSION: ${NCCL_MINOR_VERSION}")
endif()
if (${NCCL_MAJOR_VERSION} GREATER 2)
add_definitions(-DUSE_NCCL_P2P=1)
else()
if(${NCCL_MAJOR_VERSION} EQUAL 2)
if(${NCCL_MINOR_VERSION} GREATER 7)
add_definitions(-DUSE_NCCL_P2P=1)
endif()
endif()
endif()
set(NCCL_INCLUDE_DIRS ${NCCL_INCLUDE_DIR})
set(NCCL_LIBRARIES ${NCCL_LIBRARY})
message( STATUS "NCCL (include: ${NCCL_INCLUDE_DIRS}, library: ${NCCL_LIBRARIES})" )

11
orttraining/orttraining/models/runner/training_runner.cc
View file

@ -20,7 +20,9 @@
#include "orttraining/core/framework/distributed_run_context.h"
#include "orttraining/core/graph/optimizer_graph_builder.h"
#include "orttraining/models/runner/training_util.h"
#if defined(USE_NCCL) && defined(USE_NCCL_P2P)
#include "orttraining/training_ops/cuda/communication/nccl_service.h"
#endif
#include "single_include/nlohmann/json.hpp"
#include "test/perftest/utils.h"
@ -706,8 +708,8 @@ Status TrainingRunner::TrainingLoop(IDataLoader& training_data_loader, IDataLoad
auto end_to_end_start = std::chrono::high_resolution_clock::now();
bool end_to_end_measurement_started = false;
#ifdef USE_NCCL
// NCCL-P2P
#if defined(USE_NCCL) && defined(USE_NCCL_P2P)
// Create communication plan.
auto& nccl_service = cuda::NcclService::GetInstance();
nccl_service.PlanStart();
@ -727,6 +729,7 @@ Status TrainingRunner::TrainingLoop(IDataLoader& training_data_loader, IDataLoad
}
nccl_service.PlanEnd();
// Launch NCCL service to execute the plan.
nccl_service.Launch();
#endif
@ -779,7 +782,7 @@ Status TrainingRunner::TrainingLoop(IDataLoader& training_data_loader, IDataLoad
fetch_names,
fetches));
RunWithUpdate(feed_names, fetch_names, feeds, fetches);
#ifdef USE_NCCL
#if defined(USE_NCCL) && defined(USE_NCCL_P2P)
nccl_service.Reset();
#endif
} else {
@ -904,7 +907,7 @@ Status TrainingRunner::TrainingLoop(IDataLoader& training_data_loader, IDataLoad
<< "Average Step Time: " << all_steps_duration_seconds.count() / (step_ - step_start) << " Second\n"
<< "Average Step Throughput: " << params_.batch_size * (step_ - step_start) / (all_steps_duration_seconds.count()) << " Examples / Second\n";
#ifdef USE_NCCL
#if defined(USE_NCCL) && defined(USE_NCCL_P2P)
nccl_service.Terminate();
#endif
return Status::OK();

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

@ -1,7 +1,7 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#ifdef USE_NCCL
#if defined(USE_NCCL) && defined(USE_NCCL_P2P)
#include "core/common/common.h"
#include "core/profile/context.h"

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

@ -1,7 +1,7 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#ifdef USE_NCCL
#if defined(USE_NCCL) && defined(USE_NCCL_P2P)
#pragma once
#include <condition_variable>

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

@ -74,7 +74,7 @@ void Recv::ReceiveData(
recvRange.Begin();
#endif
#ifdef USE_NCCL
#if defined(USE_NCCL) && defined(USE_NCCL_P2P)
buffer = GetScratchBuffer<char>(aggregated_aligned_tensor_bytes);
#else
buffer = AllocateBufferOnCPUPinned<char>(static_cast<size_t>(aggregated_aligned_tensor_bytes));
@ -87,7 +87,7 @@ void Recv::ReceiveData(
// The following NCCL call is equivalent to the following MPI call. User can
// uncomment the MPI call to debug.
#ifdef USE_NCCL
#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
@ -118,7 +118,7 @@ void Recv::ReceiveData(
tensor_offset_in_bytes = GetAggregatedAlignedAddress(tensor_offset_in_bytes);
// Copy data out from buffer.
#ifdef USE_NCCL
#if defined(USE_NCCL) && defined(USE_NCCL_P2P)
CUDA_CALL(cudaMemcpyAsync(tensor->MutableDataRaw(), buffer.get() + tensor_offset_in_bytes,
tensor->SizeInBytes(), cudaMemcpyHostToDevice));
#else
@ -128,7 +128,8 @@ void Recv::ReceiveData(
tensor_offset_in_bytes += tensor->SizeInBytes();
}
#ifndef USE_NCCL
#if defined(USE_NCCL) && defined(USE_NCCL_P2P)
#else
AddDeferredReleaseCPUPtr(buffer.release());
#endif

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

@ -100,7 +100,7 @@ void Send::SendData(
memcpyRange.Begin();
#endif
#ifdef USE_NCCL
#if defined(USE_NCCL) && defined(USE_NCCL_P2P)
IAllocatorUniquePtr<char> buffer = GetScratchBuffer<char>(aggregated_aligned_tensor_bytes);
#else
IAllocatorUniquePtr<char> buffer = AllocateBufferOnCPUPinned<char>(
@ -109,7 +109,7 @@ void Send::SendData(
for (int i = 0; i < num_tensors; ++i) {
const Tensor* tensor = ctx->Input<Tensor>(i + 2);
#ifdef USE_NCCL
#if defined(USE_NCCL) && defined(USE_NCCL_P2P)
CUDA_CALL(cudaMemcpy(buffer.get() + tensor_offsets_in_bytes[i], tensor->DataRaw(),
tensor_sizes_in_bytes[i], cudaMemcpyDeviceToDevice));
#else
@ -137,7 +137,7 @@ void Send::SendData(
dst,
static_cast<int>(tag_)};
#ifdef USE_NCCL
#if defined(USE_NCCL) && defined(USE_NCCL_P2P)
auto& nccl_service = cuda::NcclService::GetInstance();
nccl_service.SubmitSendAndWait(info_data.buffer, info_data.size, info_data.rank);
#else