Add CPU send/recv for pipeline (#5315)

* cpu send/recv

* clean up send/recv

* remove unused code

* assert and nccl option for mnist

* add build option to enable build with only cpu. Without this, nccl is always enabled which will break build on machine that only contains cpu

* Add USE_MPI distinct from USE_NCCL/USE_HOROVOD

* fix

* fix

* exclude cpu send/recv for machines without mpi

Co-authored-by: Tim Harris <tiharr@microsoft.com>

cmake/CMakeLists.txt

@@ -1216,6 +1216,7 @@ if (onnxruntime_ENABLE_TRAINING)
     include(FindPackageHandleStandardArgs)
     find_package_handle_standard_args(MPI DEFAULT_MSG MPI_INCLUDE_DIR MPI_LIBRARY)
 
+    set(onnxruntime_USE_MPI ON)
     if (MPI_FOUND)
       if(NOT DEFINED onnxruntime_MPI_HOME)
         execute_process(COMMAND mpirun --version OUTPUT_VARIABLE MPIRUN_OUTPUT)
@@ -1238,6 +1239,7 @@ if (onnxruntime_ENABLE_TRAINING)
       endif()
 
     else ()
+      set(onnxruntime_USE_MPI OFF)
       set(onnxruntime_USE_NCCL OFF)
       set(onnxruntime_USE_HOROVOD OFF)
       message( WARNING "MPI is not found. Please use --mpi_home to specify the path of MPI. Otherwise, NCCL or HOROVOD will be disabled." )
@@ -1308,6 +1310,10 @@ if (onnxruntime_ENABLE_TRAINING)
     endif()
   endif()
 
+  if (onnxruntime_USE_MPI AND MPI_FOUND)
+    add_definitions(-DUSE_MPI=1)
+  endif()
+
   if (onnxruntime_USE_HOROVOD AND MPI_FOUND)
     if (WIN32)
       message( FATAL_ERROR "Horovod is not supported on Windows." )

cmake/onnxruntime_providers.cmake

@@ -183,7 +183,8 @@ if (onnxruntime_ENABLE_TRAINING)
   if (onnxruntime_USE_HOROVOD)
     target_include_directories(onnxruntime_providers PRIVATE ${HOROVOD_INCLUDE_DIRS})
   endif()
-  if (onnxruntime_USE_NCCL OR onnxruntime_USE_HOROVOD)
+
+  if (onnxruntime_USE_MPI)
     target_include_directories(onnxruntime_providers PUBLIC ${MPI_INCLUDE_DIRS})
   endif()
 endif()

onnxruntime/core/framework/allocation_planner.cc

@@ -476,8 +476,10 @@ class PlannerImpl {
         OrtValueIndex index = Index(node_output->Name());
         ProcessDef(index, node_output);
         ++UseCount(index);
+        auto allocator = exec_provider->GetAllocator(0, p_kernel_def->OutputMemoryType(i));
+        ORT_ENFORCE(allocator);
         plan_.SetLocation(static_cast<size_t>(index),
-                          exec_provider->GetAllocator(0, p_kernel_def->OutputMemoryType(i))->Info());
+                          allocator->Info());
       }
 
       // if sync is needed, mark allocation plan as create_fence_if_async=true

orttraining/orttraining/core/framework/mpi_context.cc

@@ -11,13 +11,13 @@ local_rank_(0),
 world_size_(1),
 local_size_(1)
 {
-#if defined(USE_NCCL)
+#if defined(USE_MPI)
   setup_mpi();
 #endif
 }
 
 MPIContext::~MPIContext() {
-#if defined(USE_NCCL)
+#if defined(USE_MPI)
 #ifndef _WIN32
   shutdown_mpi();
 #endif
@@ -29,7 +29,7 @@ const MPIContext& MPIContext::GetInstance() {
   return context;
 }
 
-#if defined(USE_NCCL)
+#if defined(USE_MPI)
 void MPIContext::setup_mpi() {
   // setup MPI amd horovod
   int is_mpi_initialized = 0;

orttraining/orttraining/core/framework/mpi_context.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#if defined(USE_NCCL)
+#if defined(USE_MPI)
 #include <mpi.h>
 #endif
 
@@ -27,7 +27,7 @@ class MPIContext {
 
     MPIContext(MPIContext const&) = delete;
     void operator=(MPIContext const&) = delete;
-    
+
     // within ~MPIContext() we need to check for _WIN32 before calling shutdown_mpi().
     ~MPIContext();
 
@@ -36,7 +36,7 @@ class MPIContext {
     int GetWorldSize() const { return world_size_; }
     int GetLocalSize() const { return local_size_; }
 
-#if defined(USE_NCCL)
+#if defined(USE_MPI)
     // https://docs.microsoft.com/en-us/windows/win32/dlls/dynamic-link-library-best-practices
     // in case of _WIN32 we cannot call shutdown_mpi() in MPIContext destructor because of DllMain's restriction
     // shutdown_mpi shall be called specifically in user code.
@@ -46,7 +46,7 @@ class MPIContext {
   private:
     MPIContext();
 
-#if defined(USE_NCCL)
+#if defined(USE_MPI)
     void setup_mpi();
 #endif
     int world_rank_;

orttraining/orttraining/core/framework/mpi_setup.h

@@ -0,0 +1,30 @@
+#pragma once
+
+#if defined(USE_MPI)
+// Disable deprecated C++ bindings, which fail to type-check with GCC 9.3
+#define OMPI_SKIP_MPICXX
+#include <mpi.h>
+#endif
+
+#ifdef USE_HOROVOD
+#include "orttraining/core/graph/horovod_adapters.h"
+#endif
+
+namespace onnxruntime {
+namespace training {
+
+struct MPIContext {
+  MPIContext(int world_rank = 0, int local_rank = 0, int world_size = 1, int local_size = 1);
+  int world_rank;
+  int local_rank;
+  int world_size;
+  int local_size;
+};
+
+#if defined(USE_MPI)
+MPIContext setup_mpi();
+void shutdown_mpi();
+#endif
+
+}  // namespace training
+}  // namespace onnxruntime

orttraining/orttraining/models/bert/main.cc

@@ -553,7 +553,7 @@ void setup_training_params(BertParameters& params) {
   params.model_with_training_graph_path = model_name_base + ORT_TSTR("_bw.onnx");
   params.model_actual_running_graph_path = model_name_base + ORT_TSTR("_bw_running.onnx");
 
-#if defined(USE_NCCL) || defined(USE_HOROVOD)
+#if defined(USE_MPI)
   if (params.pipeline_parallel_size > 1) {
     auto pipeline_model_name_base = model_name_base + ToPathString(std::to_string(MPIContext::GetInstance().GetWorldRank()));
     params.pipeline_partitioned_model_path = pipeline_model_name_base + ORT_TSTR("_partitioned.onnx");
@@ -841,7 +841,7 @@ int main(int argc, char* argv[]) {
     RETURN_IF_FAIL(RunTraining(params, *env));
   }
 
-#if defined(USE_NCCL)
+#if defined(USE_MPI)
 #ifdef _WIN32
   // https://docs.microsoft.com/en-us/windows/win32/dlls/dynamic-link-library-best-practices
   // shutdown_mpi() is not called within MPIContext destructor because of DllMain's restriction

orttraining/orttraining/models/gpt2/main.cc

@@ -297,7 +297,7 @@ void setup_training_params(GPT2Parameters& params) {
                                            {/*prediction_name*/ "output",
                                             /*label_name*/ "labels"});
 
-#if defined(USE_NCCL) || defined(USE_HOROVOD)
+#if defined(USE_MPI)
   ORT_ENFORCE(params.horizontal_parallel_size <= MPIContext::GetInstance().GetWorldSize());
   ORT_ENFORCE(params.data_parallel_size <= MPIContext::GetInstance().GetWorldSize());
   if (MPIContext::GetInstance().GetWorldSize() % params.horizontal_parallel_size != 0) {
@@ -481,7 +481,7 @@ int main(int argc, char* argv[]) {
     RETURN_IF_FAIL(RunTraining(params, *env));
   }
 
-#if defined(USE_NCCL)
+#if defined(USE_MPI)
 #ifdef _WIN32
   // https://docs.microsoft.com/en-us/windows/win32/dlls/dynamic-link-library-best-practices
   // shutdown_mpi() is not called within MPIContext destructor because of DllMain's restriction

orttraining/orttraining/models/mnist/main.cc

@@ -53,6 +53,16 @@ Status ParseArguments(int argc, char* argv[], TrainingRunner::Parameters& params
       ("train_batch_size", "Total batch size for training.", cxxopts::value<int>()->default_value("100"))
       ("eval_batch_size", "Total batch size for eval.", cxxopts::value<int>()->default_value("100"))
       ("learning_rate", "The initial learning rate for Adam.", cxxopts::value<float>()->default_value("0.01"))
+      ("display_loss_steps", "How often to dump loss into tensorboard", cxxopts::value<size_t>()->default_value("10"))
+      ("use_nccl", "Whether to use NCCL for distributed training.", cxxopts::value<bool>()->default_value("false"))
+      ("data_parallel_size", "Data parallel group size.", cxxopts::value<int>()->default_value("1"))
+      ("horizontal_parallel_size", "Horizontal model parallel group size.", cxxopts::value<int>()->default_value("1"))
+      ("pipeline_parallel_size", "Number of pipeline stages.", cxxopts::value<int>()->default_value("1"))
+      ("cut_group_info", "Specify the cutting info for graph partition (pipeline only). An example of a cut_group_info of "
+      "size two is: 1393:407-1463/1585/1707,2369:407-2439/2561/2683. Here, the cut info is split by ',', with the first "
+      "cut_info equal to 1393:407-1463/1585/1707, and second cut_info equal to 2369:407-2439/2561/2683. Each CutEdge is "
+      "seperated by ':'. If consumer nodes need to be specified, specify them after producer node with a '-' delimiter and "
+      "separate each consumer node with a '/'. ", cxxopts::value<std::vector<std::string>>()->default_value(""))
       ("evaluation_period", "How many training steps to make before making an evaluation.",
         cxxopts::value<size_t>()->default_value("1"));
   // clang-format on
@@ -80,13 +90,73 @@ Status ParseArguments(int argc, char* argv[], TrainingRunner::Parameters& params
     params.train_data_dir.assign(train_data_dir.begin(), train_data_dir.end());
     params.log_dir.assign(log_dir.begin(), log_dir.end());
     params.use_profiler = flags.count("use_profiler") > 0;
+    params.display_loss_steps = flags["display_loss_steps"].as<size_t>();
+    params.data_parallel_size = flags["data_parallel_size"].as<int>();
+    params.horizontal_parallel_size = flags["horizontal_parallel_size"].as<int>();
+    // pipeline_parallel_size controls the number of pipeline's stages.
+    // pipeline_parallel_size=1 means no model partition, which means all processes run
+    // the same model. We only partition model when pipeline_parallel_size > 1.
+    params.pipeline_parallel_size = flags["pipeline_parallel_size"].as<int>();
+    ORT_RETURN_IF_NOT(params.data_parallel_size > 0, "data_parallel_size must > 0");
+    ORT_RETURN_IF_NOT(params.horizontal_parallel_size > 0, "horizontal_parallel_size must > 0");
+    ORT_RETURN_IF_NOT(params.pipeline_parallel_size > 0, "pipeline_parallel_size must > 0");
 
+    // If user doesn't provide partitioned model files, a cut list should be provided for ORT to do partition
+    // online. If the pipeline contains n stages, the cut list should be of length (n-1), in order to cut the
+    // graph into n partitions.
+    if (params.pipeline_parallel_size > 1) {
+      auto cut_info_groups = flags["cut_group_info"].as<std::vector<std::string>>();
+
+      ORT_RETURN_IF_NOT(static_cast<int>(cut_info_groups.size() + 1) == params.pipeline_parallel_size,
+                        "cut_info length plus one must match pipeline parallel size");
+
+      auto process_with_delimiter = [](std::string& input_str, const std::string& delimiter) {
+        std::vector<std::string> result;
+        size_t pos = 0;
+        std::string token;
+        while ((pos = input_str.find(delimiter)) != std::string::npos) {
+          token = input_str.substr(0, pos);
+          result.emplace_back(token);
+          input_str.erase(0, pos + delimiter.length());
+        }
+        // push the last split of substring into result.
+        result.emplace_back(input_str);
+        return result;
+      };
+
+      auto process_cut_info = [&](std::string& cut_info_string) {
+        TrainingSession::TrainingConfiguration::CutInfo cut_info;
+        const std::string edge_delimiter = ":";
+        const std::string consumer_delimiter = "/";
+        const std::string producer_consumer_delimiter = "-";
+
+        auto cut_edges = process_with_delimiter(cut_info_string, edge_delimiter);
+        for (auto& cut_edge : cut_edges) {
+          auto process_edge = process_with_delimiter(cut_edge, producer_consumer_delimiter);
+          if (process_edge.size() == 1) {
+            TrainingSession::TrainingConfiguration::CutEdge edge{process_edge[0]};
+            cut_info.emplace_back(edge);
+          } else {
+            ORT_ENFORCE(process_edge.size() == 2);
+            auto consumer_list = process_with_delimiter(process_edge[1], consumer_delimiter);
+
+            TrainingSession::TrainingConfiguration::CutEdge edge{process_edge[0], consumer_list};
+            cut_info.emplace_back(edge);
+          }
+        }
+        return cut_info;
+      };
+
+      for (auto& cut_info : cut_info_groups) {
+        TrainingSession::TrainingConfiguration::CutInfo cut = process_cut_info(cut_info);
+        params.pipeline_partition_cut_list.emplace_back(cut);
+      }
+    }
+    params.use_nccl = flags["use_nccl"].as<bool>();
 #ifdef USE_CUDA
     bool use_cuda = flags.count("use_cuda") > 0;
     if (use_cuda) {
-      // Use local rank as device ID of the associated CUDA EP.
-      OrtDevice::DeviceId device_id = static_cast<OrtDevice::DeviceId>(MPIContext::GetInstance().GetLocalRank());
-      params.providers.emplace(kCudaExecutionProvider, CreateExecutionProviderFactory_CUDA(device_id));
+      params.providers.emplace(kCudaExecutionProvider, CreateExecutionProviderFactory_CUDA(0));
     }
 #endif
   } catch (const exception& e) {

orttraining/orttraining/python/orttraining_pybind_state.cc

@@ -154,7 +154,7 @@ TrainingConfigurationResult ConfigureSessionForTraining(
   return python_config_result;
 }
 
-#if defined(USE_NCCL)
+#if defined(USE_MPI)
 void CopyMPIContextToTrainingParameters(TrainingParameters& parameters, const logging::Logger* logger) {
   LOGS(*logger, INFO) << "MPIContext::GetInstance().GetWorldRank(): " << MPIContext::GetInstance().GetWorldRank();
   LOGS(*logger, INFO) << "MPIContext::GetInstance().GetLocalRank(): " << MPIContext::GetInstance().GetLocalRank();
@@ -203,7 +203,7 @@ void addObjectMethodsForTraining(py::module& m) {
       .def_readwrite("transformer_layer_recompute", &TrainingParameters::transformer_layer_recompute)
       .def_readwrite("number_recompute_layers", &TrainingParameters::number_recompute_layers);
 
-#if defined(USE_NCCL)
+#if defined(USE_MPI)
   m.def("get_mpi_context_local_rank", []() -> int { return MPIContext::GetInstance().GetLocalRank(); });
   m.def("get_mpi_context_local_size", []() -> int { return MPIContext::GetInstance().GetLocalSize(); });
   m.def("get_mpi_context_world_rank", []() -> int { return MPIContext::GetInstance().GetWorldRank(); });
@@ -237,7 +237,7 @@ void addObjectMethodsForTraining(py::module& m) {
         return onnxruntime::make_unique<PyTrainingSession>(env, GetDefaultCPUSessionOptions());
       }))
       .def("finalize", [](py::object) {
-#if defined(USE_NCCL)
+#if defined(USE_MPI)
 #ifdef _WIN32
         // https://docs.microsoft.com/en-us/windows/win32/dlls/dynamic-link-library-best-practices
         // shutdown_mpi() is not called within MPIContext destructor because of DllMain's restriction
@@ -249,7 +249,7 @@ void addObjectMethodsForTraining(py::module& m) {
       .def("load_model", [](PyTrainingSession* sess, const std::string& path, TrainingParameters& parameters) {
         OrtPybindThrowIfError(sess->GetSessionHandle()->Load(path));
 
-#if defined(USE_NCCL)
+#if defined(USE_MPI)
         bool use_nccl = parameters.allreduce_post_accumulation;
         if (!use_nccl && parameters.world_size > 1)
           CopyMPIContextToTrainingParameters(parameters, sess->GetSessionHandle()->GetLogger());
@@ -265,7 +265,7 @@ void addObjectMethodsForTraining(py::module& m) {
         std::istringstream buffer(serialized_model);
         OrtPybindThrowIfError(sess->GetSessionHandle()->Load(buffer));
 
-#if defined(USE_NCCL)
+#if defined(USE_MPI)
         bool use_nccl = parameters.allreduce_post_accumulation;
         if (!use_nccl && parameters.world_size > 1)
           CopyMPIContextToTrainingParameters(parameters, sess->GetSessionHandle()->GetLogger());

orttraining/orttraining/test/model/training_runner_test.cc

@@ -60,6 +60,56 @@ TEST(TrainingRunnerTest, Basic) {
   ASSERT_TRUE(status.IsOK()) << status.ErrorMessage();
 }
 
+// This test verifies a simple pipeline run with 3 partitions.
+// TODO: enable this test after distributed run is enabled in CI.
+TEST(TrainingRunnerTest, DISABLED_PipelineRun) {
+  TrainingRunner::Parameters params{};
+  params.model_path = k_original_model_path;
+  params.model_with_training_graph_path = k_backward_model_path;
+  params.is_perf_test = false;
+  params.batch_size = 1;
+  params.eval_batch_size = 1;
+  params.num_train_steps = 1;
+  params.display_loss_steps = 10;
+  params.pipeline_parallel_size = 3;
+  params.fetch_names = {"predictions"};
+
+  params.loss_func_info = LossFunctionInfo(OpDef("MeanSquaredError"), "loss", {"predictions", "labels"});
+
+  // cut model in 3 partitions
+  TrainingSession::TrainingConfiguration::CutInfo cut0 = {
+      onnxruntime::training::TrainingSession::TrainingConfiguration::CutEdge("T3")};
+
+  TrainingSession::TrainingConfiguration::CutInfo cut1 = {
+      onnxruntime::training::TrainingSession::TrainingConfiguration::CutEdge("T6")};
+
+  params.pipeline_partition_cut_list.emplace_back(cut0);
+  params.pipeline_partition_cut_list.emplace_back(cut1);
+
+  std::unique_ptr<Environment> env;
+  ASSERT_TRUE(Environment::Create(nullptr, env).IsOK());
+  TrainingRunner runner{params, *env};
+
+  auto status = runner.Initialize();
+  ASSERT_TRUE(status.IsOK()) << status.ErrorMessage();
+
+  std::vector<std::string> tensor_names{
+      "X", "labels"};
+  std::vector<TensorShape> tensor_shapes{
+      {1, 784}, {1, 10}};
+  std::vector<ONNX_NAMESPACE::TensorProto_DataType> tensor_types{
+      ONNX_NAMESPACE::TensorProto_DataType_FLOAT, ONNX_NAMESPACE::TensorProto_DataType_FLOAT};
+
+  auto data_set = std::make_shared<RandomDataSet>(1, tensor_names, tensor_shapes, tensor_types);
+  auto data_loader = std::make_shared<SingleDataLoader>(data_set, tensor_names);
+
+  status = runner.Run(data_loader.get(), data_loader.get());
+  ASSERT_TRUE(status.IsOK()) << status.ErrorMessage();
+
+  status = runner.EndTraining(data_loader.get());
+  ASSERT_TRUE(status.IsOK()) << status.ErrorMessage();
+}
+
 }  // namespace test
 }  // namespace training
 }  // namespace onnxruntime

orttraining/orttraining/training_ops/communication_common.h

@@ -2,11 +2,11 @@
 // Licensed under the MIT License.
 
 #include "core/common/common.h"
-#include "core/providers/cuda/cuda_common.h"
 #pragma once
 
+#include "orttraining/core/framework/mpi_context.h"
+
 namespace onnxruntime {
-namespace cuda {
 
 typedef struct {
   // Pointer to ONNX tensor's data on CPU.
@@ -102,5 +102,92 @@ inline void ComputeShapeRelatedInfo(
   }
 }
 
-}  // namespace cuda
+inline void SendShapeInfo(
+    const int dst,
+    const int64_t tag,      // mpi send tag
+    const int num_tensors,  // Number of sent tensors.
+    size_t aggregated_aligned_tensor_bytes,
+    std::vector<size_t> prefix_tensor_shape_sizes,
+    std::vector<int64_t> aggregated_tensor_shapes) {
+  const int num_tensors_in_bytes = num_tensors * static_cast<int>(sizeof(size_t));
+  ORT_ENFORCE(num_tensors_in_bytes < INT_MAX,
+              "Total tensor number larger than MPI size limit");
+
+  CommInfo_t info_shape_sizes{prefix_tensor_shape_sizes.data(),
+                              num_tensors_in_bytes,
+                              dst,
+                              static_cast<int>(tag)};
+  ORT_ENFORCE(aggregated_aligned_tensor_bytes < INT_MAX,
+              "Aggregated tensor size larger than MPI size limit");
+
+  CommInfo_t info_aggregated_size{&aggregated_aligned_tensor_bytes,
+                                  static_cast<int>(sizeof(size_t)),
+                                  dst,
+                                  static_cast<int>(tag)};
+
+  int total_tensor_dim_in_bytes = static_cast<int>(
+                                      aggregated_tensor_shapes.size()) *
+                                  static_cast<int>(sizeof(int64_t));
+  ORT_ENFORCE(total_tensor_dim_in_bytes < INT_MAX,
+              "Total dimensions of tensors larger than MPI size limit");
+
+  CommInfo_t info_shapes{aggregated_tensor_shapes.data(),
+                         total_tensor_dim_in_bytes,
+                         dst,
+                         static_cast<int>(tag)};
+
+  // Directly use CPU to wait MPI_Send. We cannot use GPU callback because
+  // MPI_Send may block the entire GPU until it returns.
+  MPI_CHECK(MPI_Send(
+      info_shape_sizes.buffer, info_shape_sizes.size, MPI_CHAR,
+      info_shape_sizes.rank, info_shape_sizes.tag, MPI_COMM_WORLD));
+
+  MPI_CHECK(MPI_Send(
+      info_aggregated_size.buffer, info_aggregated_size.size, MPI_CHAR,
+      info_aggregated_size.rank, info_aggregated_size.tag, MPI_COMM_WORLD));
+
+  MPI_CHECK(MPI_Send(
+      info_shapes.buffer, info_shapes.size, MPI_CHAR,
+      info_shapes.rank, info_shapes.tag, MPI_COMM_WORLD));
+}
+
+inline void ReceiveShapeInfo(
+    const int src,
+    const int64_t tag,  // mpi recv tag
+    const int num_tensors,
+    size_t& aggregated_aligned_tensor_bytes,
+    std::vector<size_t>& prefix_tensor_shape_sizes,
+    std::vector<int64_t>& aggregated_tensor_shapes) {
+  // Resize vector so that the following .data() returns meaningful pointer.
+  prefix_tensor_shape_sizes.resize(num_tensors);
+  CommInfo_t info_shape_sizes{prefix_tensor_shape_sizes.data(),
+                              num_tensors * static_cast<int>(sizeof(size_t)),
+                              src,
+                              static_cast<int>(tag)};
+  CommInfo_t info_aggregated_size{&aggregated_aligned_tensor_bytes,
+                                  static_cast<int>(sizeof(size_t)),
+                                  src,
+                                  static_cast<int>(tag)};
+  // Directly use CPU to wait MPI_Recv. We cannot use GPU callback because
+  // MPI_Recv may block the entire GPU until it returns.
+  MPI_CHECK(MPI_Recv(
+      info_shape_sizes.buffer, info_shape_sizes.size, MPI_CHAR,
+      info_shape_sizes.rank, info_shape_sizes.tag, MPI_COMM_WORLD, MPI_STATUS_IGNORE));
+
+  MPI_CHECK(MPI_Recv(
+      info_aggregated_size.buffer, info_aggregated_size.size, MPI_CHAR,
+      info_aggregated_size.rank, info_aggregated_size.tag, MPI_COMM_WORLD, MPI_STATUS_IGNORE));
+
+  // prefix_tensor_shape_sizes's last element is the number of total dimensions.
+  // If a 3-D tensor and a 2-D tensor are sent, its value is 2 + 3 = 5.
+  aggregated_tensor_shapes.resize(prefix_tensor_shape_sizes[num_tensors - 1]);
+  CommInfo_t info_shapes{aggregated_tensor_shapes.data(),
+                         static_cast<int>(aggregated_tensor_shapes.size()) * static_cast<int>(sizeof(int64_t)),
+                         src,
+                         static_cast<int>(tag)};
+  MPI_CHECK(MPI_Recv(
+      info_shapes.buffer, info_shapes.size, MPI_CHAR,
+      info_shapes.rank, info_shapes.tag, MPI_COMM_WORLD, MPI_STATUS_IGNORE));
+}
+
 }  // namespace onnxruntime

orttraining/orttraining/training_ops/cpu/communication/recv.cc

@@ -0,0 +1,170 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+#if defined(USE_MPI)
+
+#include "orttraining/training_ops/cpu/communication/recv.h"
+
+#include <mpi.h>
+
+#include "orttraining/training_ops/communication_common.h"
+#include "orttraining/core/framework/mpi_context.h"
+
+namespace onnxruntime {
+namespace contrib {
+
+void Recv::ReceiveData(
+    const int num_tensors,
+    std::vector<Tensor*> received_tensors,
+    const int src,
+    const size_t aggregated_aligned_tensor_bytes,
+    std::vector<char>& buffer) const {
+  buffer.reserve(aggregated_aligned_tensor_bytes);
+  CommInfo_t info_data{buffer.data(),
+                       static_cast<int>(aggregated_aligned_tensor_bytes),
+                       src,
+                       static_cast<int>(tag_)};
+
+  MPI_CHECK(MPI_Recv(
+      info_data.buffer, info_data.size, MPI_CHAR,
+      info_data.rank, info_data.tag, MPI_COMM_WORLD, MPI_STATUS_IGNORE));
+
+  // Copy tensors from buffer to outputs.
+  size_t tensor_offset_in_bytes = 0;
+  for (int i = 0; i < num_tensors; ++i) {
+    Tensor* tensor = received_tensors[i];
+
+    // Find the next aligned offset in the tensor buffer to meet alignment requirement
+    tensor_offset_in_bytes = GetAggregatedAlignedAddress(tensor_offset_in_bytes);
+
+    // Copy data out from buffer.
+    assert(tensor_offset_in_bytes + tensor->SizeInBytes() <= aggregated_aligned_tensor_bytes);
+    memcpy(tensor->MutableDataRaw(), buffer.data() + tensor_offset_in_bytes, tensor->SizeInBytes());
+
+    tensor_offset_in_bytes += tensor->SizeInBytes();
+  }
+  assert(tensor_offset_in_bytes == aggregated_aligned_tensor_bytes);
+}
+
+ONNX_OPERATOR_KERNEL_EX(
+    Recv,
+    kMSDomain,
+    1,
+    kCpuExecutionProvider,
+    KernelDefBuilder()
+        .InputMemoryType<OrtMemTypeDefault>(0)  /* CPU variable */
+        .InputMemoryType<OrtMemTypeDefault>(1)  /* CPU variable */
+        .OutputMemoryType<OrtMemTypeDefault>(0) /* CPU variable */
+        .TypeConstraint("TBool", DataTypeImpl::GetTensorType<bool>())
+        .TypeConstraint("TInt64", DataTypeImpl::GetTensorType<int64_t>())
+        .TypeConstraint("V", DataTypeImpl::AllFixedSizeTensorTypes()),
+    Recv);
+
+Status Recv::Compute(OpKernelContext* ctx) const {
+  // Check if control signal is true.
+  const Tensor* input_signal_tensor = ctx->Input<Tensor>(0);
+  const bool* input_signal = input_signal_tensor->template Data<bool>();
+  ORT_ENFORCE(*input_signal, "Input control signal of Recv must be true before executing the node.");
+
+  // Extract remote rank
+  const Tensor* remote_rank_tensor = ctx->Input<Tensor>(1);
+  const int64_t* remote_rank = remote_rank_tensor->template Data<int64_t>();
+  const int src = static_cast<int>(*remote_rank);
+
+  // Start communication
+  int world_rank;
+  MPI_CHECK(MPI_Comm_rank(MPI_COMM_WORLD, &world_rank));
+  ORT_ENFORCE(world_rank != src, "Receive data from rank ", src, " on the rank ", world_rank, ".");
+
+  const int num_tensors = static_cast<int>(element_types_.size());
+  std::vector<size_t> tensor_sizes_in_bytes;
+  std::vector<TensorShape> tensor_shapes;
+  // TODO move the following variables to member variables for extending life-time
+  // if we want to make the entire call async
+  size_t aggregated_aligned_tensor_bytes = 0;
+  std::vector<size_t> prefix_tensor_shape_sizes;
+  std::vector<int64_t> aggregated_tensor_shapes;
+  // tensor_offsets_in_bytes[i] is the starting byte of the i-th tensor in the send tensor buffer
+  std::vector<size_t> tensor_offsets_in_bytes;
+
+  // Whether shapes are statically inferrable.
+  bool all_shapes_inferred = true;
+  // At iteration i, the i-th received tensor is processed.
+  for (int i = 0; i < num_tensors; ++i) {
+    TensorShape inferred_shape;
+    // The first input is a boolean control signal. We only check actual received tensors.
+    auto shape_inferred = ctx->TryGetInferredOutputShape(i + 1, inferred_shape);
+    if (!shape_inferred) {
+      all_shapes_inferred = false;
+      break;
+    }
+  }
+
+  std::vector<Tensor*> received_tensors(num_tensors);
+  if (all_shapes_inferred) {
+    // Create outputs before communication because all shapes are inferred.
+    for (int i = 0; i < num_tensors; ++i) {
+      TensorShape inferred_shape;
+      // The first input is a boolean control signal. We only work on actual received tensors before that.
+      ORT_ENFORCE(ctx->TryGetInferredOutputShape(i + 1, inferred_shape));
+      // If shape is statically inferred, we declare output here and
+      // access its shape from operator's context in GetTensorShapesAndSizes(...).
+      received_tensors[i] = ctx->Output(i + 1, inferred_shape);
+    }
+
+    GetTensorShapesAndSizes(
+        false,        // value of "is_index_input". Received tensors are "output"s so this flag is "false".
+        1,            // First received tensor's index.
+        num_tensors,  // Number of tensors to received.
+        ctx,
+        tensor_sizes_in_bytes,
+        tensor_shapes);
+
+    // Extract information needed for copying input tensors from a big buffer
+    // to individual locations.
+    // Only that big buffer will be received through MPI.
+    ComputeShapeRelatedInfo(
+        tensor_sizes_in_bytes,
+        tensor_shapes,
+        aggregated_aligned_tensor_bytes,
+        prefix_tensor_shape_sizes,
+        aggregated_tensor_shapes,
+        tensor_offsets_in_bytes);
+  } else {
+    ReceiveShapeInfo(
+        src,
+        tag_,
+        num_tensors,
+        aggregated_aligned_tensor_bytes,
+        prefix_tensor_shape_sizes,
+        aggregated_tensor_shapes);
+
+    // Create output tensors. Unlike the case where we can infer output shapes before communication,
+    // we need to create outputs after receiving shapes.
+    size_t begin = 0;
+    for (int i = 0; i < num_tensors; ++i) {
+      std::vector<int64_t> tensor_shape(aggregated_tensor_shapes.begin() + begin,
+                                        aggregated_tensor_shapes.begin() + prefix_tensor_shape_sizes[i]);
+      received_tensors[i] = ctx->Output(i + 1, tensor_shape);
+      // Move the "begin" to the beginning dimension of the next received tensor.
+      begin = prefix_tensor_shape_sizes[i];
+    }
+  }
+
+  // At this stage, all shape information (either inferred locally or received from the source process)
+  // required to receive tensors are ready.
+  // Create buffer and receive data.
+  std::vector<char> buffer;
+  ReceiveData(num_tensors, received_tensors, src, aggregated_aligned_tensor_bytes, buffer);
+
+  // Set first output after communication is done.
+  Tensor* output_signal_tensor = ctx->Output(0, {});
+  bool* output_signal = output_signal_tensor->template MutableData<bool>();
+  *output_signal = true;
+
+  return Status::OK();
+}
+
+}  // namespace contrib
+}  // namespace onnxruntime
+
+#endif

orttraining/orttraining/training_ops/cpu/communication/recv.h

@@ -0,0 +1,33 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+#if defined(USE_MPI)
+#pragma once
+#include "core/common/common.h"
+#include "core/framework/op_kernel.h"
+
+namespace onnxruntime {
+namespace contrib {
+
+class Recv final : public OpKernel {
+public:
+  Recv(const OpKernelInfo& info) : OpKernel(info) {
+    ORT_ENFORCE(info.GetAttr<int64_t>("tag", &tag_).IsOK());
+    ORT_ENFORCE(info.GetAttrs<int64_t>("element_types", element_types_).IsOK());
+  }
+
+  Status Compute(OpKernelContext* context) const override;
+
+private:
+  void ReceiveData(
+    const int num_tensors,
+    std::vector<Tensor*> received_tensors,
+    const int src,
+    const size_t aggregated_aligned_tensor_bytes,
+    std::vector<char>& buffer) const;
+  int64_t tag_;
+  std::vector<int64_t> element_types_;
+};
+
+}  // namespace cuda
+}  // namespace onnxruntime
+#endif

orttraining/orttraining/training_ops/cpu/communication/send.cc

@@ -0,0 +1,126 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+#if defined(USE_MPI)
+#include "orttraining/training_ops/cpu/communication/send.h"
+
+#include <mpi.h>
+
+#include "orttraining/training_ops/communication_common.h"
+#include "orttraining/core/framework/mpi_context.h"
+
+namespace onnxruntime {
+namespace contrib {
+
+ONNX_OPERATOR_KERNEL_EX(
+    Send,
+    kMSDomain,
+    1,
+    kCpuExecutionProvider,
+    KernelDefBuilder()
+        .InputMemoryType<OrtMemTypeDefault>(0)  /* CPU variable */
+        .InputMemoryType<OrtMemTypeDefault>(1)  /* CPU variable */
+        .OutputMemoryType<OrtMemTypeDefault>(0) /* CPU variable */
+        .TypeConstraint("TBool", DataTypeImpl::GetTensorType<bool>())
+        .TypeConstraint("TInt64", DataTypeImpl::GetTensorType<int64_t>())
+        .TypeConstraint("V", DataTypeImpl::AllFixedSizeTensorTypes()),
+    Send);
+
+void Send::SendData(
+    OpKernelContext* ctx,
+    const int dst,
+    const int num_tensors,
+    size_t aggregated_aligned_tensor_bytes,
+    std::vector<size_t> tensor_offsets_in_bytes,
+    std::vector<size_t> tensor_sizes_in_bytes) const {
+  std::vector<char> buffer;
+  buffer.reserve(aggregated_aligned_tensor_bytes);
+
+  for (int i = 0; i < num_tensors; ++i) {
+    const Tensor* tensor = ctx->Input<Tensor>(i + 2);
+    memcpy(buffer.data() + tensor_offsets_in_bytes[i], tensor->DataRaw(), tensor_sizes_in_bytes[i]);
+  }
+
+  CommInfo_t info_data{buffer.data(),
+                       static_cast<int>(aggregated_aligned_tensor_bytes),
+                       dst,
+                       static_cast<int>(tag_)};
+
+  MPI_CHECK(MPI_Send(
+      info_data.buffer, info_data.size, MPI_CHAR,
+      info_data.rank, info_data.tag, MPI_COMM_WORLD));
+}
+
+Status Send::Compute(OpKernelContext* ctx) const {
+  // Check if control signal is true.
+  const Tensor* input_signal_tensor = ctx->Input<Tensor>(0);
+  const bool* input_signal = input_signal_tensor->template Data<bool>();
+  ORT_ENFORCE(*input_signal, "Input control signal of Send must be true before executing the node.");
+
+  // Extract remote rank
+  const Tensor* remote_rank_tensor = ctx->Input<Tensor>(1);
+  const int64_t* remote_rank = remote_rank_tensor->template Data<int64_t>();
+  const int dst = static_cast<int>(*remote_rank);
+
+  // Same-rank communication is not allowed because we currently don't have async Send/Recv.
+  int world_rank;
+  MPI_CHECK(MPI_Comm_rank(MPI_COMM_WORLD, &world_rank));
+  ORT_ENFORCE(world_rank != dst, "Sending data to rank ", dst, " on the rank ", world_rank, ".");
+
+  const int num_tensors = static_cast<int>(element_types_.size());
+  std::vector<size_t> tensor_sizes_in_bytes;
+  std::vector<TensorShape> tensor_shapes;
+  GetTensorShapesAndSizes(
+      true,
+      2,            // First sent tensor's index.
+      num_tensors,  // Number of tensors to send
+      ctx,
+      tensor_sizes_in_bytes,
+      tensor_shapes);
+
+  // TODO move the following variables to member variables for extending life-time
+  // if we want to make the entire call async
+  size_t aggregated_aligned_tensor_bytes = 0;
+  std::vector<size_t> prefix_tensor_shape_sizes;
+  std::vector<int64_t> aggregated_tensor_shapes;
+  // tensor_offsets_in_bytes[i] is the starting byte of the i-th tensor in the send tensor buffer
+  std::vector<size_t> tensor_offsets_in_bytes;
+
+  // Extract information needed for copying input tensors into a big buffer.
+  // Only that big buffer will be sent.
+  ComputeShapeRelatedInfo(
+      tensor_sizes_in_bytes,
+      tensor_shapes,
+      aggregated_aligned_tensor_bytes,
+      prefix_tensor_shape_sizes,
+      aggregated_tensor_shapes,
+      tensor_offsets_in_bytes);
+
+  bool all_shapes_inferred = true;
+  for (int i = 0; i < num_tensors; ++i) {
+    TensorShape inferred_shape;
+    auto shape_inferred = ctx->TryGetInferredInputShape(i + 2, inferred_shape);
+    if (!shape_inferred) {
+      all_shapes_inferred = false;
+      break;
+    }
+  }
+
+  // Communicate shape information when it cannot be inferred.
+  if (!all_shapes_inferred) {
+    SendShapeInfo(dst, tag_, num_tensors, aggregated_aligned_tensor_bytes, prefix_tensor_shape_sizes, aggregated_tensor_shapes);
+  }
+
+  // Send tensors.
+  SendData(ctx, dst, num_tensors, aggregated_aligned_tensor_bytes, tensor_offsets_in_bytes, tensor_sizes_in_bytes);
+
+  // Communication is done, so output control signal can be set to true.
+  Tensor* output_signal_tensor = ctx->Output(0, {});
+  bool* output_signal = output_signal_tensor->MutableData<bool>();
+  *output_signal = true;
+
+  return Status::OK();
+}
+
+}  // namespace contrib
+}  // namespace onnxruntime
+#endif

orttraining/orttraining/training_ops/cpu/communication/send.h

@@ -0,0 +1,35 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+#if defined(USE_MPI)
+#pragma once
+#include "core/common/common.h"
+#include "core/framework/op_kernel.h"
+
+namespace onnxruntime {
+namespace contrib {
+
+class Send final : public OpKernel {
+ public:
+  Send(const OpKernelInfo& info) : OpKernel(info) {
+    ORT_ENFORCE(info.GetAttr<int64_t>("tag", &tag_).IsOK());
+    ORT_ENFORCE(info.GetAttrs<int64_t>("element_types", element_types_).IsOK());
+  }
+
+  Status Compute(OpKernelContext* context) const override;
+
+ private:
+  void SendData(
+      OpKernelContext* ctx,
+      const int dst,
+      const int num_tensors,
+      size_t aggregated_aligned_tensor_bytes,
+      std::vector<size_t> tensor_offsets_in_bytes,
+      std::vector<size_t> tensor_sizes_in_bytes) const;
+
+  int64_t tag_;
+  std::vector<int64_t> element_types_;
+};
+
+}  // namespace contrib
+}  // namespace onnxruntime
+#endif

orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc

@@ -104,6 +104,12 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1,
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double_int64_t, Scale);
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double_int32_t, Scale);
 
+#ifdef USE_MPI
+// Pipeline communication operators.
+class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Send);
+class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Recv);
+#endif
+
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, RecordEvent);
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, WaitEvent);
 
@@ -205,6 +211,11 @@ Status RegisterCpuTrainingKernels(KernelRegistry& kernel_registry) {
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double_int64_t, Scale)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double_int32_t, Scale)>,
 
+#ifdef USE_MPI
+      BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Send)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Recv)>,
+#endif
+
       BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, RecordEvent)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, WaitEvent)>};
 

orttraining/orttraining/training_ops/cuda/communication/recv.cc

@@ -4,12 +4,11 @@
 #if defined(USE_NCCL) || defined(USE_HOROVOD)
 
 #include "orttraining/training_ops/cuda/communication/recv.h"
-#include "orttraining/training_ops/cuda/communication/common.h"
+#include "orttraining/training_ops/communication_common.h"
 #include "orttraining/training_ops/cuda/communication/nccl_service.h"
 #include "core/profile/profile.h"
 #include "core/profile/context.h"
 #include "core/providers/cuda/cuda_common.h"
-#include <string>
 #include <mpi.h>
 
 #include "orttraining/core/framework/mpi_context.h"
@@ -17,44 +16,6 @@
 namespace onnxruntime {
 namespace cuda {
 
-void Recv::ReceiveShapeInfo(
-    const int src,
-    const int num_tensors,
-    size_t& aggregated_aligned_tensor_bytes,
-    std::vector<size_t>& prefix_tensor_shape_sizes,
-    std::vector<int64_t>& aggregated_tensor_shapes) const {
-  // Resize vector so that the following .data() returns meaningful pointer.
-  prefix_tensor_shape_sizes.resize(num_tensors);
-  CommInfo_t info_shape_sizes{prefix_tensor_shape_sizes.data(),
-                              num_tensors * static_cast<int>(sizeof(size_t)),
-                              src,
-                              static_cast<int>(tag_)};
-  CommInfo_t info_aggregated_size{&aggregated_aligned_tensor_bytes,
-                                  static_cast<int>(sizeof(size_t)),
-                                  src,
-                                  static_cast<int>(tag_)};
-  // Directly use CPU to wait MPI_Recv. We cannot use GPU callback because
-  // MPI_Recv may block the entire GPU until it returns.
-  MPI_CHECK(MPI_Recv(
-      info_shape_sizes.buffer, info_shape_sizes.size, MPI_CHAR,
-      info_shape_sizes.rank, info_shape_sizes.tag, MPI_COMM_WORLD, MPI_STATUS_IGNORE));
-
-  MPI_CHECK(MPI_Recv(
-      info_aggregated_size.buffer, info_aggregated_size.size, MPI_CHAR,
-      info_aggregated_size.rank, info_aggregated_size.tag, MPI_COMM_WORLD, MPI_STATUS_IGNORE));
-
-  // prefix_tensor_shape_sizes's last element is the number of total dimensions.
-  // If a 3-D tensor and a 2-D tensor are sent, its value is 2 + 3 = 5.
-  aggregated_tensor_shapes.resize(prefix_tensor_shape_sizes[num_tensors - 1]);
-  CommInfo_t info_shapes{aggregated_tensor_shapes.data(),
-                         static_cast<int>(aggregated_tensor_shapes.size()) * static_cast<int>(sizeof(int64_t)),
-                         src,
-                         static_cast<int>(tag_)};
-  MPI_CHECK(MPI_Recv(
-      info_shapes.buffer, info_shapes.size, MPI_CHAR,
-      info_shapes.rank, info_shapes.tag, MPI_COMM_WORLD, MPI_STATUS_IGNORE));
-}
-
 void Recv::ReceiveData(
     const int num_tensors,
     std::vector<Tensor*> received_tensors,
@@ -119,6 +80,7 @@ void Recv::ReceiveData(
     // Find the next aligned offset in the tensor buffer to meet alignment requirement
     tensor_offset_in_bytes = GetAggregatedAlignedAddress(tensor_offset_in_bytes);
 
+    assert(tensor_offset_in_bytes + tensor->SizeInBytes() <= aggregated_aligned_tensor_bytes);
     // Copy data out from buffer.
 #if defined(USE_NCCL) && defined(USE_NCCL_P2P)
     CUDA_CALL(cudaMemcpy(tensor->MutableDataRaw(), buffer.get() + tensor_offset_in_bytes,
@@ -129,6 +91,7 @@ void Recv::ReceiveData(
 #endif
     tensor_offset_in_bytes += tensor->SizeInBytes();
   }
+  assert(tensor_offset_in_bytes == aggregated_aligned_tensor_bytes);
 
 #if defined(USE_NCCL) && defined(USE_NCCL_P2P)
 #else
@@ -240,6 +203,7 @@ Status Recv::ComputeInternal(OpKernelContext* ctx) const {
   } else {
     ReceiveShapeInfo(
         src,
+        tag_,
         num_tensors,
         aggregated_aligned_tensor_bytes,
         prefix_tensor_shape_sizes,

orttraining/orttraining/training_ops/cuda/communication/recv.h

@@ -20,12 +20,6 @@ public:
   Status ComputeInternal(OpKernelContext* context) const override;
 
 private:
-  void ReceiveShapeInfo(
-      const int src,
-      const int num_tensors,
-      size_t& aggregated_aligned_tensor_bytes,
-      std::vector<size_t>& prefix_tensor_shape_sizes,
-      std::vector<int64_t>& aggregated_tensor_shapes) const;
   void ReceiveData(
     const int num_tensors,
     std::vector<Tensor*> received_tensors,

orttraining/orttraining/training_ops/cuda/communication/send.cc

@@ -4,13 +4,11 @@
 #if defined(USE_NCCL) || defined(USE_HOROVOD)
 
 #include "orttraining/training_ops/cuda/communication/send.h"
-#include "orttraining/training_ops/cuda/communication/common.h"
+#include "orttraining/training_ops/communication_common.h"
 #include "orttraining/training_ops/cuda/communication/nccl_service.h"
 #include "core/profile/profile.h"
 #include "core/profile/context.h"
 #include "core/providers/cuda/cuda_common.h"
-#include <limits>
-#include <string>
 #include <mpi.h>
 
 #include "orttraining/core/framework/mpi_context.h"
@@ -32,53 +30,6 @@ ONNX_OPERATOR_KERNEL_EX(
         .TypeConstraint("V", DataTypeImpl::AllFixedSizeTensorTypes()),
     Send);
 
-void Send::SendShapeInfo(
-    const int dst,
-    const int num_tensors,  // Number of sent tensors.
-    size_t aggregated_aligned_tensor_bytes,
-    std::vector<size_t> prefix_tensor_shape_sizes,
-    std::vector<int64_t> aggregated_tensor_shapes) const {
-  const int num_tensors_in_bytes = num_tensors * static_cast<int>(sizeof(size_t));
-  ORT_ENFORCE(num_tensors_in_bytes < INT_MAX,
-              "Total tensor number larger than MPI size limit");
-
-  CommInfo_t info_shape_sizes{prefix_tensor_shape_sizes.data(),
-                              num_tensors_in_bytes,
-                              dst,
-                              static_cast<int>(tag_)};
-  ORT_ENFORCE(aggregated_aligned_tensor_bytes < INT_MAX,
-              "Aggregated tensor size larger than MPI size limit");
-
-  CommInfo_t info_aggregated_size{&aggregated_aligned_tensor_bytes,
-                                  static_cast<int>(sizeof(size_t)),
-                                  dst,
-                                  static_cast<int>(tag_)};
-
-  int total_tensor_dim_in_bytes = static_cast<int>(
-                                      aggregated_tensor_shapes.size()) *
-                                  static_cast<int>(sizeof(int64_t));
-  ORT_ENFORCE(total_tensor_dim_in_bytes < INT_MAX,
-              "Total dimensions of tensors larger than MPI size limit");
-
-  CommInfo_t info_shapes{aggregated_tensor_shapes.data(),
-                         total_tensor_dim_in_bytes,
-                         dst,
-                         static_cast<int>(tag_)};
-
-  // Directly use CPU to wait MPI_Send. We cannot use GPU callback because
-  // MPI_Send may block the entire GPU until it returns.
-  MPI_CHECK(MPI_Send(
-      info_shape_sizes.buffer, info_shape_sizes.size, MPI_CHAR,
-      info_shape_sizes.rank, info_shape_sizes.tag, MPI_COMM_WORLD));
-
-  MPI_CHECK(MPI_Send(
-      info_aggregated_size.buffer, info_aggregated_size.size, MPI_CHAR,
-      info_aggregated_size.rank, info_aggregated_size.tag, MPI_COMM_WORLD));
-
-  MPI_CHECK(MPI_Send(
-      info_shapes.buffer, info_shapes.size, MPI_CHAR,
-      info_shapes.rank, info_shapes.tag, MPI_COMM_WORLD));
-}
 
 void Send::SendData(
     OpKernelContext* ctx,
@@ -224,7 +175,7 @@ Status Send::ComputeInternal(OpKernelContext* ctx) const {
 
   // Communicate shape information when it cannot be inferred.
   if (!all_shapes_inferred) {
-    SendShapeInfo(dst, num_tensors, aggregated_aligned_tensor_bytes, prefix_tensor_shape_sizes, aggregated_tensor_shapes);
+    SendShapeInfo(dst, tag_, num_tensors, aggregated_aligned_tensor_bytes, prefix_tensor_shape_sizes, aggregated_tensor_shapes);
   }
 #ifdef ENABLE_NVTX_PROFILE
   // End of data preparation and shape communication.

orttraining/orttraining/training_ops/cuda/communication/send.h

@@ -21,12 +21,6 @@ public:
   Status ComputeInternal(OpKernelContext* context) const override;
 
 private:
-  void SendShapeInfo(
-      const int dst,
-      const int num_tensors,
-      size_t aggregated_aligned_tensor_bytes,
-      std::vector<size_t> prefix_tensor_shape_sizes,
-      std::vector<int64_t> aggregated_tensor_shapes) const;
   void SendData(
       OpKernelContext* ctx,
       const int dst,

tools/ci_build/build.py

@@ -97,6 +97,8 @@ def parse_arguments():
         help="Enable the pipeline c++ e2e tests.")
     parser.add_argument(
         "--use_horovod", action='store_true', help="Enable Horovod.")
+    parser.add_argument(
+        "--disable_nccl", action='store_true', help="Disable Nccl.")
     parser.add_argument(
         "--mpi_home", help="Path to MPI installation dir")
     parser.add_argument(
@@ -728,6 +730,8 @@ def generate_build_tree(cmake_path, source_dir, build_dir, cuda_home, cudnn_home
             "ON" if args.enable_training else "OFF"),
         "-Donnxruntime_USE_HOROVOD=" + (
             "ON" if args.use_horovod else "OFF"),
+        "-Donnxruntime_USE_NCCL=" + (
+            "OFF" if args.disable_nccl else "ON"),
         "-Donnxruntime_BUILD_BENCHMARKS=" + (
             "ON" if args.build_micro_benchmarks else "OFF"),
         "-Donnxruntime_USE_ROCM=" + ("ON" if args.use_rocm else "OFF"),