Remove unused NCCL subgroup support

Author: RAMitchell

src/collective/comm.cu

@@ -84,14 +84,7 @@ NCCLComm::NCCLComm(Context const* ctx, Comm const& root, std::shared_ptr<Coll> p
   rc = std::move(rc) << [&] {
     return GetUniqueId(root, this->stub_, pimpl, &nccl_unique_id_);
   } << [&] {
-    ncclConfig_t config = NCCL_CONFIG_INITIALIZER;
-    // Split subcommunicators inherit behavior from the parent communicator. With a non-blocking
-    // parent, `ncclCommSplit`/subgroup setup was unstable for the quantile tree path and failed
-    // with NCCL invalid-argument errors in MGPU tests. Keep the root communicator blocking until
-    // NCCL subgroup initialization is reliable without it.
-    config.blocking = 1;
-    return this->stub_->CommInitRankConfig(&nccl_comm_, root.World(), nccl_unique_id_, root.Rank(),
-                                           &config);
+    return this->stub_->CommInitRank(&nccl_comm_, root.World(), nccl_unique_id_, root.Rank());
   } << [&] {
     return BusyWait(this->stub_, this->nccl_comm_, this->Timeout());
   };
@@ -103,57 +96,7 @@ NCCLComm::NCCLComm(Context const* ctx, Comm const& root, std::shared_ptr<Coll> p
   }
 }
 
-std::size_t NCCLComm::GetSubgroup(std::vector<std::int32_t> const& active_ranks) const {
-  auto it = std::find_if(this->subgroups_.cbegin(), this->subgroups_.cend(),
-                         [&](auto const& g) { return g.active_ranks == active_ranks; });
-  if (it != this->subgroups_.cend()) {
-    return static_cast<std::size_t>(std::distance(this->subgroups_.cbegin(), it));
-  }
-
-  Subgroup g;
-  g.active_ranks = active_ranks;
-  g.stream = std::make_shared<curt::Stream>();
-
-  auto active =
-      std::find(active_ranks.cbegin(), active_ranks.cend(), this->Rank()) != active_ranks.cend();
-  auto color = active ? 1 : NCCL_SPLIT_NOCOLOR;
-  ncclConfig_t config = NCCL_CONFIG_INITIALIZER;
-  config.blocking = 0;
-  config.splitShare = 0;
-
-  auto rc = this->stub_->CommSplit(this->nccl_comm_, color, this->Rank(), &g.comm, &config);
-  SafeColl(rc);
-  if (g.comm != nullptr) {
-    rc = BusyWait(this->stub_, g.comm, this->Timeout());
-    SafeColl(rc);
-  }
-
-  this->subgroups_.push_back(std::move(g));
-  return this->subgroups_.size() - 1;
-}
-
 NCCLComm::~NCCLComm() {
-  for (auto& subgroup : this->subgroups_) {
-    if (subgroup.comm == nullptr) {
-      continue;
-    }
-    auto rc = Success() << [this, &subgroup] {
-      return this->stub_->CommFinalize(subgroup.comm);
-    } << [this, &subgroup] {
-      auto rc = BusyWait(this->stub_, subgroup.comm, this->Timeout());
-      if (!rc.OK()) {
-        return std::move(rc) + this->stub_->CommAbort(subgroup.comm);
-      }
-      return rc;
-    } << [this, &subgroup] {
-      return this->stub_->CommDestroy(subgroup.comm);
-    };
-    if (!rc.OK()) {
-      LOG(WARNING) << rc.Report();
-    }
-  }
-  this->subgroups_.clear();
-
   if (nccl_comm_) {
     auto rc = Success() << [this] {
       return this->stub_->CommFinalize(this->nccl_comm_);

src/collective/comm.cuh

@@ -10,7 +10,6 @@
 #include <cstdint>  // for int32_t
 #include <memory>   // for shared_ptr
 #include <utility>  // for move
-#include <vector>   // for vector
 
 #include "../common/cuda_stream.h"  // for StreamRef
 #include "coll.h"
@@ -30,20 +29,12 @@ inline Result GetCUDAResult(cudaError rc) {
 
 #if defined(XGBOOST_USE_NCCL)
 class NCCLComm : public Comm {
- public:
-  struct Subgroup {
-    std::vector<std::int32_t> active_ranks;
-    ncclComm_t comm{nullptr};
-    std::shared_ptr<curt::Stream> stream;
-  };
-
  private:
   ncclComm_t nccl_comm_{nullptr};
   std::shared_ptr<NcclStub> stub_;
   ncclUniqueId nccl_unique_id_{};
   curt::StreamRef stream_;
   std::string nccl_path_;
-  mutable std::vector<Subgroup> subgroups_;
 
  public:
   [[nodiscard]] ncclComm_t Handle() const { return nccl_comm_; }
@@ -66,8 +57,6 @@ class NCCLComm : public Comm {
     this->ResetState();
     return Success();
   }
-  [[nodiscard]] std::size_t GetSubgroup(std::vector<std::int32_t> const& active_ranks) const;
-  [[nodiscard]] Subgroup const& SubgroupAt(std::size_t idx) const { return subgroups_.at(idx); }
 };
 
 class NCCLChannel : public Channel {

src/collective/nccl_stub.cc

@@ -94,12 +94,10 @@ no long bundles NCCL in the binary wheel.
   broadcast_ = safe_load(broadcast_, "ncclBroadcast");
   allgather_ = safe_load(allgather_, "ncclAllGather");
   comm_init_rank_ = safe_load(comm_init_rank_, "ncclCommInitRank");
-  comm_init_rank_config_ = safe_load(comm_init_rank_config_, "ncclCommInitRankConfig");
   comm_destroy_ = safe_load(comm_destroy_, "ncclCommDestroy");
   comm_finalize_ = safe_load(comm_finalize_, "ncclCommFinalize");
   comm_get_async_error_ = safe_load(comm_get_async_error_, "ncclCommGetAsyncError");
   comm_abort_ = safe_load(comm_abort_, "ncclCommAbort");
-  comm_split_ = safe_load(comm_split_, "ncclCommSplit");
   get_uniqueid_ = safe_load(get_uniqueid_, "ncclGetUniqueId");
   send_ = safe_load(send_, "ncclSend");
   recv_ = safe_load(recv_, "ncclRecv");
@@ -112,12 +110,10 @@ no long bundles NCCL in the binary wheel.
   broadcast_ = ncclBroadcast;
   allgather_ = ncclAllGather;
   comm_init_rank_ = ncclCommInitRank;
-  comm_init_rank_config_ = ncclCommInitRankConfig;
   comm_destroy_ = ncclCommDestroy;
   comm_finalize_ = ncclCommFinalize;
   comm_get_async_error_ = ncclCommGetAsyncError;
   comm_abort_ = ncclCommAbort;
-  comm_split_ = ncclCommSplit;
   get_uniqueid_ = ncclGetUniqueId;
   send_ = ncclSend;
   recv_ = ncclRecv;

src/collective/nccl_stub.h

@@ -28,12 +28,10 @@ class NcclStub {
   decltype(ncclBroadcast)* broadcast_{nullptr};
   decltype(ncclAllGather)* allgather_{nullptr};
   decltype(ncclCommInitRank)* comm_init_rank_{nullptr};
-  decltype(ncclCommInitRankConfig)* comm_init_rank_config_{nullptr};
   decltype(ncclCommDestroy)* comm_destroy_{nullptr};
   decltype(ncclCommFinalize)* comm_finalize_{nullptr};
   decltype(ncclCommGetAsyncError)* comm_get_async_error_{nullptr};
   decltype(ncclCommAbort)* comm_abort_{nullptr};
-  decltype(ncclCommSplit)* comm_split_{nullptr};
   decltype(ncclGetUniqueId)* get_uniqueid_{nullptr};
   decltype(ncclSend)* send_{nullptr};
   decltype(ncclRecv)* recv_{nullptr};
@@ -68,10 +66,6 @@ class NcclStub {
                                     int rank) const {
     return this->GetNcclResult(this->comm_init_rank_(comm, nranks, commId, rank));
   }
-  [[nodiscard]] Result CommInitRankConfig(ncclComm_t* comm, int nranks, ncclUniqueId commId,
-                                          int rank, ncclConfig_t* config) const {
-    return this->GetNcclResult(this->comm_init_rank_config_(comm, nranks, commId, rank, config));
-  }
   [[nodiscard]] Result CommDestroy(ncclComm_t comm) const {
     if (this->Aborted()) {
       return Success();
@@ -100,10 +94,6 @@ class NcclStub {
     this->aborted_ = true;
     return this->GetNcclResult(comm_abort_(comm));
   }
-  [[nodiscard]] Result CommSplit(ncclComm_t comm, int color, int key, ncclComm_t* newcomm,
-                                 ncclConfig_t* config) const {
-    return this->GetNcclResult(comm_split_(comm, color, key, newcomm, config));
-  }
   [[nodiscard]] Result GetUniqueId(ncclUniqueId* uniqueId) const {
     return this->GetNcclResult(get_uniqueid_(uniqueId));
   }