From 4ad3a41ff90833634b271f17e771d7b76359d945 Mon Sep 17 00:00:00 2001 From: Ryan Hill Date: Tue, 11 May 2021 22:49:21 -0700 Subject: [PATCH] Patch through INcclService --- cmake/CMakeLists.txt | 12 +++- .../providers/cuda/cuda_provider_factory.h | 9 +++ .../providers/cuda/cuda_provider_factory.cc | 15 ++++- .../core/session/training_session.cc | 6 +- .../models/runner/training_runner.cc | 12 +++- .../cuda/communication/nccl_service.cc | 6 +- .../cuda/communication/nccl_service.h | 64 +++++++++++++++---- 7 files changed, 100 insertions(+), 24 deletions(-) diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index 1d2d583850..06b6f4106c 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -1599,9 +1599,17 @@ if (onnxruntime_ENABLE_TRAINING) endif() endif() - if (onnxruntime_USE_MPI AND MPI_FOUND) + # if (onnxruntime_USE_MPI AND MPI_FOUND) + set(MPI_INCLUDE_DIRS "C:/code/openmpi-4.0.4/ompi/include") + set(NCCL_INCLUDE_DIRS "C:/Code/nccl-1.3.4-1/src") + add_definitions(-DORT_USE_NCCL=1) + set(onnxruntime_USE_NCCL ON) add_definitions(-DUSE_MPI=1) - endif() +# endif() + +# if (onnxruntime_USE_MPI AND MPI_FOUND) + #add_definitions(-DUSE_MPI=1) + #endif() add_subdirectory(tensorboard EXCLUDE_FROM_ALL) list(APPEND onnxruntime_EXTERNAL_LIBRARIES tensorboard) diff --git a/include/onnxruntime/core/providers/cuda/cuda_provider_factory.h b/include/onnxruntime/core/providers/cuda/cuda_provider_factory.h index f321937abb..b561e5d26f 100644 --- a/include/onnxruntime/core/providers/cuda/cuda_provider_factory.h +++ b/include/onnxruntime/core/providers/cuda/cuda_provider_factory.h @@ -13,6 +13,11 @@ struct IExecutionProviderFactory; struct CUDAExecutionProviderInfo; enum class ArenaExtendStrategy : int32_t; struct CUDAExecutionProviderExternalAllocatorInfo; + +namespace cuda { +class INcclService; +} + } // namespace onnxruntime struct ProviderInfo_CUDA { @@ -35,6 +40,10 @@ struct ProviderInfo_CUDA { virtual int cudaGetDeviceCount() = 0; virtual void CUDAExecutionProviderInfo__FromProviderOptions(const onnxruntime::ProviderOptions& options, onnxruntime::CUDAExecutionProviderInfo& info) = 0; +#if defined(USE_CUDA) && defined(ORT_USE_NCCL) && defined(USE_NCCL_P2P) + virtual cuda::INcclService& GetINcclService() = 0; +#endif + virtual std::shared_ptr CreateExecutionProviderFactory(const onnxruntime::CUDAExecutionProviderInfo& info) = 0; virtual std::shared_ptr CreateCudaAllocator(int16_t device_id, size_t gpu_mem_limit, onnxruntime::ArenaExtendStrategy arena_extend_strategy, onnxruntime::CUDAExecutionProviderExternalAllocatorInfo& external_allocator_info, OrtArenaCfg* default_memory_arena_cfg) = 0; }; diff --git a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc index cb776bc43f..20906f349c 100644 --- a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc +++ b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc @@ -18,6 +18,12 @@ using namespace onnxruntime; namespace onnxruntime { +#if defined(USE_CUDA) && defined(ORT_USE_NCCL) && defined(USE_NCCL_P2P) +namespace cuda { +cuda::INcclService& GetINcclService(); +} +#endif + void Shutdown_DeleteRegistry(); struct CUDAProviderFactory : IExecutionProviderFactory { @@ -39,9 +45,6 @@ std::shared_ptr CreateExecutionProviderFactory_CUDA(c return std::make_shared(info); } -} // namespace onnxruntime - -namespace onnxruntime { struct ProviderInfo_CUDA_Impl : ProviderInfo_CUDA { OrtStatus* SetCurrentGpuDeviceId(_In_ int device_id) override { int num_devices; @@ -129,6 +132,12 @@ struct ProviderInfo_CUDA_Impl : ProviderInfo_CUDA { info = CUDAExecutionProviderInfo::FromProviderOptions(options); } +#if defined(USE_CUDA) && defined(ORT_USE_NCCL) && defined(USE_NCCL_P2P) + cuda::INcclService& GetINcclService() override { + return cuda::GetINcclService(); + } +#endif + std::shared_ptr CreateExecutionProviderFactory(const CUDAExecutionProviderInfo& info) override { return std::make_shared(info); } diff --git a/orttraining/orttraining/core/session/training_session.cc b/orttraining/orttraining/core/session/training_session.cc index 30130720f7..bee1f4c95f 100644 --- a/orttraining/orttraining/core/session/training_session.cc +++ b/orttraining/orttraining/core/session/training_session.cc @@ -1420,7 +1420,7 @@ std::unordered_set TrainingSession::GetTrainableModelInitializers( // Send and Recv call SubmitSendAndWait and SubmitRecvAndWait, respectively. void PipelineTrainingSession::LaunchNcclService(const int pipeline_stage_id) { ORT_ENFORCE(pipeline_stage_id >= 0, "Pipeline stage ID cannot be negative."); - auto& nccl_service = cuda::NcclService::GetInstance(); + auto& nccl_service = cuda::INcclService::GetInstance(); // Create NCCL communication plan. The plan is a vector of communication task group. // Each communication task group contains tasks which should be done in parallel. @@ -1897,7 +1897,7 @@ common::Status PipelineTrainingSession::RunWithPipeline(const RunOptions& run_op pipeline_worker_pool_.JoinAll(); onnxruntime::contrib::OrtEventPool::GetInstance().ResetAllEvents(); #if defined(USE_CUDA) && defined(ORT_USE_NCCL) && defined(USE_NCCL_P2P) - auto& nccl_service = cuda::NcclService::GetInstance(); + auto& nccl_service = cuda::INcclService::GetInstance(); nccl_service.Reset(); #endif @@ -1906,7 +1906,7 @@ common::Status PipelineTrainingSession::RunWithPipeline(const RunOptions& run_op PipelineTrainingSession::~PipelineTrainingSession() { #if defined(USE_CUDA) && defined(ORT_USE_NCCL) && defined(USE_NCCL_P2P) - auto& nccl_service = cuda::NcclService::GetInstance(); + auto& nccl_service = cuda::INcclService::GetInstance(); nccl_service.Terminate(); #endif } diff --git a/orttraining/orttraining/models/runner/training_runner.cc b/orttraining/orttraining/models/runner/training_runner.cc index ce3cf7cc79..b64745cc9d 100644 --- a/orttraining/orttraining/models/runner/training_runner.cc +++ b/orttraining/orttraining/models/runner/training_runner.cc @@ -23,6 +23,16 @@ #include "orttraining/training_ops/cpu/controlflow/event_pool.h" #if defined(USE_CUDA) && defined(ORT_USE_NCCL) && defined(USE_NCCL_P2P) #include "orttraining/training_ops/cuda/communication/nccl_service.h" +#include "core/providers/cuda/cuda_provider_factory.h" +namespace onnxruntime { +ProviderInfo_CUDA* GetProviderInfo_CUDA(); + +namespace cuda { +INcclService& INcclService::GetInstance() { + return GetProviderInfo_CUDA()->GetINcclService(); +} +} // namespace cuda +} // namespace onnxruntime #endif #include "nlohmann/json.hpp" #include "test/perftest/utils.h" @@ -764,7 +774,7 @@ Status TrainingRunner::TrainingLoop(IDataLoader& training_data_loader, IDataLoad #if defined(USE_CUDA) && defined(ORT_USE_NCCL) && defined(USE_NCCL_P2P) // Create communication plan. - auto& nccl_service = cuda::NcclService::GetInstance(); + auto& nccl_service = cuda::INcclService::GetInstance(); nccl_service.PlanStart(); for (auto& slot : pipeline_schedule_.GetSchedule(pipeline_context_.pipeline_stage_id)) { diff --git a/orttraining/orttraining/training_ops/cuda/communication/nccl_service.cc b/orttraining/orttraining/training_ops/cuda/communication/nccl_service.cc index 7118affb1e..b320441603 100644 --- a/orttraining/orttraining/training_ops/cuda/communication/nccl_service.cc +++ b/orttraining/orttraining/training_ops/cuda/communication/nccl_service.cc @@ -13,6 +13,10 @@ namespace onnxruntime { namespace cuda { +INcclService& GetINcclService { + return NcclService::GetInstance(); +} + bool NcclTask::Compare(const NcclTask& other) const { if (type != other.type) { return false; @@ -375,4 +379,4 @@ void NcclService::Terminate() { } // namespace cuda } // namespace onnxruntime -#endif \ No newline at end of file +#endif diff --git a/orttraining/orttraining/training_ops/cuda/communication/nccl_service.h b/orttraining/orttraining/training_ops/cuda/communication/nccl_service.h index b108357cce..48a48fde60 100644 --- a/orttraining/orttraining/training_ops/cuda/communication/nccl_service.h +++ b/orttraining/orttraining/training_ops/cuda/communication/nccl_service.h @@ -81,7 +81,7 @@ struct NcclTaskGroup final { // Below is an example of planning tasks. Notice that the communication operations in the same group are // called in random order, so those operations cannot have mutual dependency. // -// auto& nccl_service = cuda::NcclService::GetInstance(); +// auto& nccl_service = cuda::INcclService::GetInstance(); // // nccl_service.PlanStart(); // Signal the begin of communication planning. // @@ -96,7 +96,43 @@ struct NcclTaskGroup final { // nccl_service.PlanEndNewGroup(); // Mark the end of the second time slot. // // nccl_service.EndPlan(); // Signal the end of communication planning. -class NcclService final { +class INcclService { + public: + static INcclService& GetInstance(); + + // Planning APIs. They are not thread-safe. + + // Mark the start of entire plan. + virtual void PlanStart() = 0; + // Mark the end of entire plan. + virtual void PlanEnd() = 0; + // Mark the begin of a new communication group. It uses the latest time slot. + // Operations in a group can happen in random order. + virtual void PlanNewGroupStart() = 0; + // Mark the end of the current communication group. + virtual void PlanNewGroupEnd() = 0; + // Add Send to the current communication group. + virtual void PlanSend(const int dst) = 0; + // Add Recv to the current communication group. + virtual void PlanRecv(const int src) = 0; + + // Runtime APIs. They are thread-safe. + + // Launch NCCL service. It's an infinite loop which repeatedly calls corresponding NCCL + // when planned operators (e.g., Send and Recv) arrive. + virtual void Launch() = 0; + // Submit a Send request with needed information such as tensor's address and number bytes to send. + virtual void SubmitSendAndWait(void* buffer, size_t count, int peer) = 0; + // Submit a Recv request with needed information such as tensor's address and number bytes to recv. + virtual void SubmitRecvAndWait(void* buffer, size_t count, int peer) = 0; + // Reset communication plan's status so that we can reuse the same communication plan for multiple + // model update steps. + virtual void Reset() = 0; + // Terminate NCCL service. + virtual void Terminate() = 0; +}; + +class NcclService final : public INcclService { public: // Get the singleton of this class. static NcclService& GetInstance() { @@ -107,33 +143,33 @@ class NcclService final { // Planning APIs. They are not thread-safe. // Mark the start of entire plan. - void PlanStart(); + void PlanStart() override; // Mark the end of entire plan. - void PlanEnd(); + void PlanEnd() override; // Mark the begin of a new communication group. It uses the latest time slot. // Operations in a group can happen in random order. - void PlanNewGroupStart(); + void PlanNewGroupStart() override; // Mark the end of the current communication group. - void PlanNewGroupEnd(); + void PlanNewGroupEnd() override; // Add Send to the current communication group. - void PlanSend(const int dst); + void PlanSend(const int dst) override; // Add Recv to the current communication group. - void PlanRecv(const int src); + void PlanRecv(const int src) override; // Runtime APIs. They are thread-safe. // Launch NCCL service. It's an infinite loop which repeatedly calls corresponding NCCL // when planned operators (e.g., Send and Recv) arrive. - void Launch(); + void Launch() override; // Submit a Send request with needed information such as tensor's address and number bytes to send. - void SubmitSendAndWait(void* buffer, size_t count, int peer); + void SubmitSendAndWait(void* buffer, size_t count, int peer) override; // Submit a Recv request with needed information such as tensor's address and number bytes to recv. - void SubmitRecvAndWait(void* buffer, size_t count, int peer); + void SubmitRecvAndWait(void* buffer, size_t count, int peer) override; // Reset communication plan's status so that we can reuse the same communication plan for multiple // model update steps. - void Reset(); + void Reset() override; // Terminate NCCL service. - void Terminate(); + void Terminate() override; // Print debug string. friend std::ostream& operator<<(std::ostream& stream, const NcclService& service); @@ -184,4 +220,4 @@ class NcclService final { } // namespace cuda } // namespace onnxruntime -#endif \ No newline at end of file +#endif