[REVIEW] Enable Multi-Node Multi-GPU functionality

src/common/device_helpers.cuh

@@ -23,6 +23,10 @@
 
 #ifdef XGBOOST_USE_NCCL
 #include "nccl.h"
+#include "../common/io.h"
+#else
+#define NCCL_UNIQUE_ID_BYTES 128
+typedef struct { char internal[NCCL_UNIQUE_ID_BYTES]; } ncclUniqueId;
 #endif
 
 // Uncomment to enable
@@ -853,6 +857,7 @@ class AllReducer {
   std::vector<ncclComm_t> comms;
   std::vector<cudaStream_t> streams;
   std::vector<int> device_ordinals;  // device id from CUDA
+  bool use_nccl_opg{false}; // NCCL single process per GPU
 #endif
 
  public:
@@ -866,16 +871,31 @@ class AllReducer {
    * group.
    *
    * \param device_ordinals The device ordinals.
+   * \param opg Whether to initialize AllReducer with a single process per GPU
    */
 
-  void Init(const std::vector<int> &device_ordinals) {
+  void Init(const std::vector<int> &device_ordinals, bool opg=false) {
 #ifdef XGBOOST_USE_NCCL
     /** \brief this >monitor . init. */
     this->device_ordinals = device_ordinals;
+    this->use_nccl_opg = opg;
     comms.resize(device_ordinals.size());
-    dh::safe_nccl(ncclCommInitAll(comms.data(),
-                                  static_cast<int>(device_ordinals.size()),
-                                  device_ordinals.data()));
+
+    if (use_nccl_opg) {
+      CHECK_EQ(device_ordinals.size(), 1)
+        << "NCCL-OPG version currently supports only 1 GPU in a process.";
+      auto id = GetUniqueId();
+      dh::safe_nccl(ncclCommInitRank(
+      	&(comms[0]),
+      	rabit::GetWorldSize(),
+      	id, rabit::GetRank()));
+    } else {
+      dh::safe_nccl(ncclCommInitAll(
+      	comms.data(),
+      	static_cast<int>(device_ordinals.size()),
+      	device_ordinals.data()));
+    }
+
     streams.resize(device_ordinals.size());
     for (size_t i = 0; i < device_ordinals.size(); i++) {
       safe_cuda(cudaSetDevice(device_ordinals[i]));
@@ -910,7 +930,9 @@ class AllReducer {
    */
   void GroupStart() {
 #ifdef XGBOOST_USE_NCCL
-    dh::safe_nccl(ncclGroupStart());
+    if (!use_nccl_opg) {
+      dh::safe_nccl(ncclGroupStart());
+    }
 #endif
   }
 
@@ -919,7 +941,9 @@ class AllReducer {
    */
   void GroupEnd() {
 #ifdef XGBOOST_USE_NCCL
-    dh::safe_nccl(ncclGroupEnd());
+	if (!use_nccl_opg) {
+	  dh::safe_nccl(ncclGroupEnd());
+	}
 #endif
   }
 
@@ -1009,6 +1033,29 @@ class AllReducer {
       dh::safe_cuda(cudaSetDevice(device_ordinals[i]));
       dh::safe_cuda(cudaStreamSynchronize(streams[i]));
     }
+  #endif
+  };
+
+  /**
+   * \fn  ncclUniqueId GetUniqueId()
+   *
+   * \brief Gets the Unique ID from NCCL to be used in setting up interprocess
+   * communication
+   *
+   * \return the Unique ID
+   */
+  ncclUniqueId GetUniqueId() {
+#ifdef XGBOOST_USE_NCCL
+    static const int RootRank = 0;
+    ncclUniqueId id;
+    if (rabit::GetRank() == RootRank) {
+      dh::safe_nccl(ncclGetUniqueId(&id));
+    }
+    rabit::Broadcast(
+      (void*)&id,
+      (size_t)sizeof(ncclUniqueId),
+      (int)RootRank);
+    return id;
 #endif
   }
 };

src/tree/param.h

@@ -76,6 +76,8 @@ struct TrainParam : public dmlc::Parameter<TrainParam> {
   int gpu_id;
   // number of GPUs to use
   int n_gpus;
+  // enable single process per GPU in NCCL
+  bool distributed;
   // the criteria to use for ranking splits
   std::string split_evaluator;
 
@@ -199,6 +201,9 @@ struct TrainParam : public dmlc::Parameter<TrainParam> {
         .set_lower_bound(-1)
         .set_default(1)
         .describe("Number of GPUs to use for multi-gpu algorithms: -1=use all GPUs");
+    DMLC_DECLARE_FIELD(distributed)
+        .set_default(false)
+        .describe("EXP Param: Set this to True to enable a single process per GPU.");
     DMLC_DECLARE_FIELD(split_evaluator)
         .set_default("elastic_net,monotonic,interaction")
         .describe("The criteria to use for ranking splits");

src/tree/updater_gpu_hist.cu

@@ -992,7 +992,7 @@ class GPUHistMakerSpecialised{
       device_list_[index] = device_id;
     }
 
-    reducer_.Init(device_list_);
+    reducer_.Init(device_list_, param_.distributed);
 
     auto batch_iter = dmat->GetRowBatches().begin();
     const SparsePage& batch = *batch_iter;
@@ -1050,6 +1050,7 @@ class GPUHistMakerSpecialised{
 
   void AllReduceHist(int nidx) {
     if (shards_.size() == 1) return;
+    dh::safe_cuda(cudaDeviceSynchronize());
     monitor_.Start("AllReduce");
 
     reducer_.GroupStart();
@@ -1080,6 +1081,9 @@ class GPUHistMakerSpecialised{
         right_node_max_elements, shard->ridx_segments[nidx_right].Size());
     }
 
+    rabit::Allreduce<rabit::op::Max, size_t>(&left_node_max_elements, 1);
+    rabit::Allreduce<rabit::op::Max, size_t>(&right_node_max_elements, 1);
+
     auto build_hist_nidx = nidx_left;
     auto subtraction_trick_nidx = nidx_right;
 
@@ -1142,6 +1146,9 @@ class GPUHistMakerSpecialised{
           tmp_sums[i] = dh::SumReduction(
               shard->temp_memory, shard->gpair.Data(), shard->gpair.Size());
         });
+
+    rabit::Allreduce<rabit::op::Sum>((GradientPair::ValueT*)&tmp_sums[0], 2);
+
     GradientPair sum_gradient =
         std::accumulate(tmp_sums.begin(), tmp_sums.end(), GradientPair());