diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index d2041c099d..119e013016 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -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})" ) diff --git a/orttraining/orttraining/models/runner/training_runner.cc b/orttraining/orttraining/models/runner/training_runner.cc index 084d74858b..bb1e406493 100644 --- a/orttraining/orttraining/models/runner/training_runner.cc +++ b/orttraining/orttraining/models/runner/training_runner.cc @@ -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(); diff --git a/orttraining/orttraining/training_ops/cuda/communication/nccl_service.cc b/orttraining/orttraining/training_ops/cuda/communication/nccl_service.cc index f85719be0f..d67c093505 100644 --- a/orttraining/orttraining/training_ops/cuda/communication/nccl_service.cc +++ b/orttraining/orttraining/training_ops/cuda/communication/nccl_service.cc @@ -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" diff --git a/orttraining/orttraining/training_ops/cuda/communication/nccl_service.h b/orttraining/orttraining/training_ops/cuda/communication/nccl_service.h index 915197bf4f..0c219d53c8 100644 --- a/orttraining/orttraining/training_ops/cuda/communication/nccl_service.h +++ b/orttraining/orttraining/training_ops/cuda/communication/nccl_service.h @@ -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 diff --git a/orttraining/orttraining/training_ops/cuda/communication/recv.cc b/orttraining/orttraining/training_ops/cuda/communication/recv.cc index 540eac031d..0758694216 100644 --- a/orttraining/orttraining/training_ops/cuda/communication/recv.cc +++ b/orttraining/orttraining/training_ops/cuda/communication/recv.cc @@ -74,7 +74,7 @@ void Recv::ReceiveData( recvRange.Begin(); #endif -#ifdef USE_NCCL +#if defined(USE_NCCL) && defined(USE_NCCL_P2P) buffer = GetScratchBuffer(aggregated_aligned_tensor_bytes); #else buffer = AllocateBufferOnCPUPinned(static_cast(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 diff --git a/orttraining/orttraining/training_ops/cuda/communication/send.cc b/orttraining/orttraining/training_ops/cuda/communication/send.cc index 25bc01460e..9b9f3b3fe7 100644 --- a/orttraining/orttraining/training_ops/cuda/communication/send.cc +++ b/orttraining/orttraining/training_ops/cuda/communication/send.cc @@ -100,7 +100,7 @@ void Send::SendData( memcpyRange.Begin(); #endif -#ifdef USE_NCCL +#if defined(USE_NCCL) && defined(USE_NCCL_P2P) IAllocatorUniquePtr buffer = GetScratchBuffer(aggregated_aligned_tensor_bytes); #else IAllocatorUniquePtr buffer = AllocateBufferOnCPUPinned( @@ -109,7 +109,7 @@ void Send::SendData( for (int i = 0; i < num_tensors; ++i) { const Tensor* tensor = ctx->Input(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(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