Remove MPI dependency (#17624)

### Description
<!-- Describe your changes. -->

Support launch multi-GPU without MPI


### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->

cmake/CMakeLists.txt

@@ -1476,29 +1476,32 @@ if (onnxruntime_ENABLE_TRAINING)
   list(APPEND onnxruntime_EXTERNAL_LIBRARIES tensorboard)
 endif()
 
-if (UNIX AND onnxruntime_USE_MPI)
-  if (EXISTS "${onnxruntime_MPI_HOME}")
-    set(MPI_HOME "${onnxruntime_MPI_HOME}")
-  elseif (EXISTS "/bert_ort/openmpi")
-    set(MPI_HOME "/bert_ort/openmpi")
+if (UNIX AND onnxruntime_USE_NCCL)
+  # MPI is INDEPENDENT of NCCL for now. You can build NCLL without MPI and launch multi-GPU with your own launcher.
+  if (onnxruntime_USE_MPI)
+    if (EXISTS "${onnxruntime_MPI_HOME}")
+      set(MPI_HOME "${onnxruntime_MPI_HOME}")
+    elseif (EXISTS "/bert_ort/openmpi")
+      set(MPI_HOME "/bert_ort/openmpi")
+    endif()
+    find_package(MPI)
+
+    if (MPI_CXX_FOUND)
+      message( STATUS "MPI Version: ${MPI_CXX_VERSION}")
+      message( STATUS "MPI (include: ${MPI_CXX_INCLUDE_DIRS}, library: ${MPI_CXX_LIBRARIES})" )
+      mark_as_advanced(MPI_CXX_INCLUDE_DIRS MPI_CXX_LIBRARIES)
+      list(APPEND onnxruntime_EXTERNAL_LIBRARIES ${MPI_CXX_LIBRARIES} ${MPI_CXX_LINK_FLAGS})
+    else ()
+      message(
+            FATAL_ERROR
+            "MPI is not found. Please define onnxruntime_MPI_HOME to specify the path of MPI. Otherwise, NCCL will be disabled."
+            "or you can remove --use_mpi from build args to disable MPI."
+          )
+    endif()
   endif()
 
-  find_package(MPI)
-
-  if (MPI_CXX_FOUND)
-    message( STATUS "MPI Version: ${MPI_CXX_VERSION}")
-    message( STATUS "MPI (include: ${MPI_CXX_INCLUDE_DIRS}, library: ${MPI_CXX_LIBRARIES})" )
-    mark_as_advanced(MPI_CXX_INCLUDE_DIRS MPI_CXX_LIBRARIES)
-    list(APPEND onnxruntime_EXTERNAL_LIBRARIES ${MPI_CXX_LIBRARIES} ${MPI_CXX_LINK_FLAGS})
-  else ()
-    message(
-          FATAL_ERROR
-          "MPI is not found. Please define onnxruntime_MPI_HOME to specify the path of MPI. Otherwise, NCCL will be disabled."
-        )
-  endif()
-
-  # Find NCCL and MPI
-  if (onnxruntime_USE_NCCL AND MPI_CXX_FOUND)
+  # Find NCCL
+  if (onnxruntime_USE_NCCL)
     if (onnxruntime_USE_CUDA)
       set(NCCL_LIBNAME "nccl")
     elseif (onnxruntime_USE_ROCM)

onnxruntime/contrib_ops/cuda/collective/nccl_kernels.cc

@@ -1,5 +1,15 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
+#include <unistd.h>
+#include <time.h>
+#include <sys/socket.h>
+#include <netinet/in.h>
+#include <arpa/inet.h>
+#include <netinet/tcp.h>
+#include <netdb.h>
+#include <cstdint>
+#include <memory>
+#include <string>
 
 #include "nccl_kernels.h"
 #include "mpi_include.h"
@@ -8,6 +18,7 @@
 #include "core/providers/cuda/math/matmul.h"
 #include "core/providers/cuda/tensor/transpose.h"
 #include "core/providers/cuda/cuda_check_memory.h"
+#include "core/platform/env_var_utils.h"
 
 namespace onnxruntime {
 namespace contrib {
@@ -38,21 +49,160 @@ static ncclDataType_t GetNcclDataType(onnxruntime::MLDataType type) {
   }
 }
 
-#ifdef USE_MPI
-static Status CreateNcclCommByMPI(int world_size, int rank, ncclComm_t* comm) {
+namespace IPC {
+#define FLLOG LOGS_DEFAULT(VERBOSE)
+#define FLLOGERRNO LOGS_DEFAULT(WARNING) << "error:" << strerror(errno)
+#define FLLOGGAI LOGS_DEFAULT(WARNING) << "error:" << gai_strerror(ret)
+
+typedef std::shared_ptr<struct addrinfo> AddrInfoPtr;
+
+int CreateSocket(bool is_server) {
+  int sockfd = -1;
+
+  struct addrinfo hints;
+  struct addrinfo* result = nullptr;
+  AddrInfoPtr result_ptr(result, [](struct addrinfo* p) { if(p){freeaddrinfo(p);} });
+
+  memset(&hints, 0, sizeof(struct addrinfo));
+  hints.ai_family = AF_UNSPEC;     /* Allow IPv4 or IPv6 */
+  hints.ai_socktype = SOCK_STREAM; /* TCP socket. use SOCK_DGRAM for UDP */
+  hints.ai_flags = AI_PASSIVE;     /* For wildcard IP address */
+  hints.ai_protocol = 0;           /* Any protocol */
+
+  std::string rank0_ip = ParseEnvironmentVariableWithDefault<std::string>("RANK0_IP", "localhost");
+  std::string port_number = ParseEnvironmentVariableWithDefault<std::string>("RANK0_PORT", "18888");
+
+  int ret = getaddrinfo(is_server ? nullptr : rank0_ip.c_str(), port_number.c_str(), &hints, &result);
+  if (ret != 0) {
+    FLLOGGAI << " getaddrinfo failed\n";
+    return sockfd;
+  }
+
+  for (struct addrinfo* rp = result; rp != nullptr; rp = rp->ai_next) {
+    sockfd = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol);
+    if (sockfd == -1) {
+      continue;
+    }
+
+    int on = 1;
+    int rc = setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, (char*)&on, sizeof(on));
+    if (rc < 0) {
+      FLLOGERRNO << ("setsockopt() failed\n");
+      close(sockfd);
+      continue;
+    }
+
+    if (is_server) {
+      if (bind(sockfd, rp->ai_addr, rp->ai_addrlen) == 0) {
+        FLLOG << "Listening on port " << port_number << " for the other GPU processores...\n";
+      } else {
+        FLLOGERRNO << ("bind failed\n");
+        close(sockfd);
+        sockfd = -1;
+      }
+    } else {
+      time_t start_time = time(0);
+      int conn_ret = connect(sockfd, rp->ai_addr, rp->ai_addrlen);
+      while (time(0) - start_time < 40 && conn_ret < 0) {
+        FLLOGERRNO << (" waiting the RANK 0 ready...\n"); /* terminate */
+        sleep(1);
+        conn_ret = connect(sockfd, rp->ai_addr, rp->ai_addrlen);
+      }
+      if (conn_ret < 0) {
+        close(sockfd);
+        sockfd = -1;
+        FLLOGERRNO << ("connect failed with timeout\n"); /* terminate */
+      } else {
+        FLLOG << "connect to " << rank0_ip << ":" << port_number << "success \n";
+      }
+    }
+    break;
+  }
+  return sockfd;
+}
+
+int WriteOnRank0(ncclUniqueId* nccl_id, int word_size) {
+  int fd = CreateSocket(true);
+  if (fd < 0) {
+    FLLOGERRNO << (" create socket\n"); /* terminate */
+    return -1;
+  }
+
+  /* listen to the socket */
+  if (listen(fd, word_size) < 0) {
+    FLLOGERRNO << ("listen\n"); /* terminate */
+    return -1;
+  }
+
+  word_size--;  // rank 0 is not in word_size
+  while (word_size-- > 0) {
+    int client_fd = accept(fd, nullptr, nullptr); /* accept blocks */
+    if (client_fd < 0) {
+      FLLOGERRNO << ("accept\n"); /* terminate */
+      return -1;
+    }
+    FLLOG << ("Accepted new GPU\n");
+    if (write(client_fd, (nccl_id), sizeof(ncclUniqueId)) != sizeof(ncclUniqueId)) {
+      FLLOGERRNO << ("write\n"); /* terminate */
+      return -1;
+    }
+    close(client_fd);
+  }
+  close(fd);
+  return 0;
+}
+
+int ReadFromRank0(ncclUniqueId* nccl_id) {
+  int sockfd = CreateSocket(false);
+  if (sockfd < 0) {
+    FLLOGERRNO << ("socket");
+    return -1;
+  }
+
+  if (read(sockfd, (nccl_id), sizeof(ncclUniqueId)) != sizeof(ncclUniqueId)) {
+    FLLOGERRNO << ("read"); /* terminate */
+    return -1;
+  }
+
+  close(sockfd);
+  return 0;
+}
+
+int IPC_Bcast(ncclUniqueId* nccl_id, int rank, int world_size) {
+  if (rank == 0) {
+    if (WriteOnRank0(nccl_id, world_size) != 0) {
+      return (-1);
+    }
+  } else if (ReadFromRank0(nccl_id) != 0) {
+    return (-1);
+  }
+
+  return 0;
+}
+}  // namespace IPC
+
+static Status CreateNcclCommunicator(int world_size, int rank, ncclComm_t* comm, bool is_launched_by_mpi) {
   // Create new NCCL communicator
   ncclUniqueId nccl_id;
   if (rank == 0) {
     NCCL_RETURN_IF_ERROR(ncclGetUniqueId(&nccl_id));
   }
-  MPI_CHECK(MPI_Bcast(&nccl_id, sizeof(nccl_id), MPI_BYTE, 0, MPI_COMM_WORLD));
+  if (is_launched_by_mpi) {
+#ifdef USE_MPI
+    MPI_CHECK(MPI_Bcast(&nccl_id, sizeof(nccl_id), MPI_BYTE, 0, MPI_COMM_WORLD));
+#else
+    return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Please compile ORT with USE_MPI.");
+#endif
+  } else if (IPC::IPC_Bcast(&nccl_id, rank, world_size) != 0) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "IPC_Bcast nccl_id failed with :", strerror(errno));
+  }
   NCCL_RETURN_IF_ERROR(ncclCommInitRank(comm, world_size, nccl_id, rank));
 
   return Status::OK();
 }
-#endif
 
 NcclContext::NcclContext() {
+  world_size_ = -1;
 #ifdef USE_MPI
   int is_mpi_initialized = 0;
   MPI_Initialized(&is_mpi_initialized);
@@ -65,14 +215,19 @@ NcclContext::NcclContext() {
   MPI_Comm_size(MPI_COMM_WORLD, &world_size_);
 
   MPI_Comm_rank(MPI_COMM_WORLD, &rank_);
+#endif
+  // world_size_ would be zero if MPI is being compiled but not launched by MPI.
+  bool is_launched_by_mpi = true;
+  if (world_size_ < 1) {
+    is_launched_by_mpi = false;
+    world_size_ = ParseEnvironmentVariableWithDefault<int32_t>("LOCAL_WORLD_SIZE", -1);
+    rank_ = ParseEnvironmentVariableWithDefault<int32_t>("LOCAL_RANK", -1);
+    ORT_ENFORCE(world_size_ != -1 && rank_ != -1);
+  }
 
   // Initialize global Parallel Group NCCL Communicator
-  auto ret = CreateNcclCommByMPI(world_size_, rank_, &comm_);
+  auto ret = CreateNcclCommunicator(world_size_, rank_, &comm_, is_launched_by_mpi);
   ORT_ENFORCE(ret.IsOK());
-
-#else
-  ORT_THROW("ORT must be built with MPI to use NCCL.");
-#endif
 }
 
 NcclContext::~NcclContext() {

onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc

@@ -145,7 +145,7 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kPytorchAtenDomain
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, ShrunkenGather);
 #endif
 
-#if defined(USE_MPI) && defined(ORT_USE_NCCL)
+#if defined(ORT_USE_NCCL)
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, AllReduce);
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, AllGather);
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, AllToAll);
@@ -301,7 +301,7 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) {
     BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, ShrunkenGather)>,
 #endif
 
-#if defined(USE_MPI) && defined(ORT_USE_NCCL)
+#if defined(ORT_USE_NCCL)
     BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, AllReduce)>,
     BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, AllGather)>,
     BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, AllToAll)>,

onnxruntime/core/graph/contrib_ops/contrib_defs.cc

@@ -3003,7 +3003,7 @@ Having this op allows runtime to do operator re-ordering to reduce compute FLOPs
   }
 #endif
 
-#ifdef USE_MPI
+#ifdef ORT_USE_NCCL
   RegisterCollectiveOps();
 #endif
 }

onnxruntime/core/graph/contrib_ops/contrib_defs.h

@@ -53,7 +53,7 @@ void RegisterContribSchemas();
 void RegisterNchwcSchemas();
 void RegisterQuantizationSchemas();
 
-#if defined(USE_MPI)
+#if defined(ORT_USE_NCCL)
 void RegisterCollectiveOps();
 #endif