[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,8 @@ class AllReducer {
   std::vector<ncclComm_t> comms;
   std::vector<cudaStream_t> streams;
   std::vector<int> device_ordinals;  // device id from CUDA
+  std::vector<int> device_counts;  // device count from CUDA
+  ncclUniqueId id;
 #endif
 
  public:
@@ -872,14 +878,43 @@ class AllReducer {
 #ifdef XGBOOST_USE_NCCL
     /** \brief this >monitor . init. */
     this->device_ordinals = device_ordinals;
-    comms.resize(device_ordinals.size());
-    dh::safe_nccl(ncclCommInitAll(comms.data(),
-                                  static_cast<int>(device_ordinals.size()),
-                                  device_ordinals.data()));
-    streams.resize(device_ordinals.size());
+    this->device_counts.resize(rabit::GetWorldSize());
+    this->comms.resize(device_ordinals.size());
+    this->streams.resize(device_ordinals.size());
+    this->id = GetUniqueId();
+
+    device_counts.at(rabit::GetRank()) = device_ordinals.size();
+    for (size_t i = 0; i < device_counts.size(); i++) {
+      rabit::Broadcast(
+        (void*)&(device_counts.at(rabit::GetRank())),
+        (size_t)sizeof(device_counts.at(rabit::GetRank())),
+        (int)rabit::GetRank());
+    }
+
+    int nccl_rank = 0;
+    int nccl_rank_offset = std::accumulate(device_counts.begin(),
+                             device_counts.begin() + rabit::GetRank(), 0);
+    int nccl_nranks = std::accumulate(device_counts.begin(),
+                        device_counts.end(), 0);
+    nccl_rank += nccl_rank_offset;
+
+    GroupStart();
     for (size_t i = 0; i < device_ordinals.size(); i++) {
-      safe_cuda(cudaSetDevice(device_ordinals[i]));
-      safe_cuda(cudaStreamCreate(&streams[i]));
+      int dev = device_ordinals.at(i);
+
+      dh::safe_cuda(cudaSetDevice(dev));
+      dh::safe_nccl(ncclCommInitRank(
+        &(comms.at(i)),
+        nccl_nranks, id, 
+        nccl_rank));
+
+      nccl_rank++;
+    }
+    GroupEnd();
+
+    for (size_t i = 0; i < device_ordinals.size(); i++) {
+      safe_cuda(cudaSetDevice(device_ordinals.at(i)));
+      safe_cuda(cudaStreamCreate(&(streams.at(i))));
     }
     initialised_ = true;
 #else
@@ -1009,6 +1044,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/updater_gpu_hist.cu

@@ -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());