Patch through INcclService

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)

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<onnxruntime::IExecutionProviderFactory> CreateExecutionProviderFactory(const onnxruntime::CUDAExecutionProviderInfo& info) = 0;
   virtual std::shared_ptr<onnxruntime::IAllocator> 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;
 };

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<IExecutionProviderFactory> CreateExecutionProviderFactory_CUDA(c
   return std::make_shared<onnxruntime::CUDAProviderFactory>(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<IExecutionProviderFactory> CreateExecutionProviderFactory(const CUDAExecutionProviderInfo& info) override {
     return std::make_shared<CUDAProviderFactory>(info);
   }

orttraining/orttraining/core/session/training_session.cc

@@ -1420,7 +1420,7 @@ std::unordered_set<std::string> 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
 }

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)) {

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
+#endif

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
+#endif