[EM] Support mixed GPU models on SNMG. (dmlc#11998)

- Query C2C info via UUID.
- Drop the specialization for full C2C to support large GPUs.

Author: trivialfis

doc/tutorials/external_memory.rst

@@ -332,11 +332,9 @@ Starting with 3.1, XGBoost introduces an adaptive cache for GPU-based external m
 training. The feature helps split the data cache into a host cache and a device cache. By
 keeping a portion of the cache on the GPU, we can reduce the amount of data transfer
 during training when there's sufficient amount of GPU memory. The feature can be
-controlled by the ``cache_host_ratio`` parameter in the
-:py:class:`xgboost.ExtMemQuantileDMatrix`. It is disabled when the device has full C2C
-bandwidth since it's not needed there. On devices that with reduced bandwidth or devices
-with PCIe connections, unless explicitly specified, the ratio is automatically estimated
-based on device memory size and the size of the dataset.
+controlled by the ``cache_host_ratio`` parameter of the
+:py:class:`xgboost.ExtMemQuantileDMatrix`. Unless explicitly specified, the ratio is
+automatically estimated based on device memory size and the size of the dataset.
 
 However, this parameter increases memory fragmentation as XGBoost needs large memory pages
 with irregular sizes. As a result, you might see out of memory error after the

src/collective/comm.cu

@@ -1,5 +1,5 @@
 /**
- * Copyright 2023-2025, XGBoost Contributors
+ * Copyright 2023-2026, XGBoost Contributors
  */
 #if defined(XGBOOST_USE_NCCL)
 #include <algorithm>  // for sort
@@ -10,15 +10,14 @@
 #include <sstream>    // for stringstream
 #include <vector>     // for vector
 
-#include "../common/cuda_context.cuh"    // for CUDAContext
-#include "../common/cuda_rt_utils.h"     // for SetDevice
-#include "../common/device_helpers.cuh"  // for DefaultStream
-#include "../common/type.h"              // for EraseType
-#include "comm.cuh"                      // for NCCLComm
-#include "comm.h"                        // for Comm
-#include "nccl_stub.h"                   // for NcclStub
-#include "xgboost/collective/result.h"   // for Result
-#include "xgboost/span.h"                // for Span
+#include "../common/cuda_context.cuh"   // for CUDAContext
+#include "../common/cuda_rt_utils.h"    // for SetDevice, GetUuid, PrintUuid
+#include "../common/type.h"             // for EraseType
+#include "comm.cuh"                     // for NCCLComm
+#include "comm.h"                       // for Comm
+#include "nccl_stub.h"                  // for NcclStub
+#include "xgboost/collective/result.h"  // for Result
+#include "xgboost/span.h"               // for Span
 
 namespace xgboost::collective {
 namespace {
@@ -38,23 +37,6 @@ Result GetUniqueId(Comm const& comm, std::shared_ptr<NcclStub> stub, std::shared
   *pid = id;
   return Success();
 }
-
-inline constexpr std::size_t kUuidLength =
-    sizeof(std::declval<cudaDeviceProp>().uuid) / sizeof(std::uint64_t);
-
-void GetCudaUUID(xgboost::common::Span<std::uint64_t, kUuidLength> const& uuid, DeviceOrd device) {
-  cudaDeviceProp prob{};
-  dh::safe_cuda(cudaGetDeviceProperties(&prob, device.ordinal));
-  std::memcpy(uuid.data(), static_cast<void*>(&(prob.uuid)), sizeof(prob.uuid));
-}
-
-std::string PrintUUID(xgboost::common::Span<std::uint64_t, kUuidLength> const& uuid) {
-  std::stringstream ss;
-  for (auto v : uuid) {
-    ss << std::hex << v;
-  }
-  return ss.str();
-}
 }  // namespace
 
 Comm* RabitComm::MakeCUDAVar(Context const* ctx, std::shared_ptr<Coll> pimpl) const {
@@ -76,18 +58,18 @@ NCCLComm::NCCLComm(Context const* ctx, Comm const& root, std::shared_ptr<Coll> p
   curt::SetDevice(ctx->Ordinal());
   stub_ = std::make_shared<NcclStub>(nccl_path);
 
-  std::vector<std::uint64_t> uuids(root.World() * kUuidLength, 0);
-  auto s_uuid = xgboost::common::Span<std::uint64_t>{uuids.data(), uuids.size()};
-  auto s_this_uuid = s_uuid.subspan(root.Rank() * kUuidLength, kUuidLength);
-  GetCudaUUID(s_this_uuid, ctx->Device());
+  std::vector<unsigned char> uuids(root.World() * curt::kUuidLength, 0);
+  auto s_uuid = common::Span{uuids.data(), uuids.size()};
+  auto s_this_uuid = s_uuid.subspan(root.Rank() * curt::kUuidLength, curt::kUuidLength);
+  curt::GetUuid(s_this_uuid, ctx->Ordinal());
 
   auto rc = pimpl->Allgather(root, common::EraseType(s_uuid));
   SafeColl(rc);
 
-  std::vector<xgboost::common::Span<std::uint64_t, kUuidLength>> converted(root.World());
+  std::vector<common::Span<unsigned char>> converted(root.World());
   std::size_t j = 0;
-  for (size_t i = 0; i < uuids.size(); i += kUuidLength) {
-    converted[j] = s_uuid.subspan(i, kUuidLength);
+  for (size_t i = 0; i < uuids.size(); i += curt::kUuidLength) {
+    converted[j] = s_uuid.subspan(i, curt::kUuidLength);
     j++;
   }
 
@@ -97,7 +79,7 @@ NCCLComm::NCCLComm(Context const* ctx, Comm const& root, std::shared_ptr<Coll> p
 
   CHECK_EQ(n_uniques, root.World())
       << "Multiple processes within communication group running on same CUDA "
-      << "device is not supported. " << PrintUUID(s_this_uuid) << "\n";
+      << "device is not supported. " << curt::PrintUuid(s_this_uuid) << "\n";
 
   rc = std::move(rc) << [&] {
     return GetUniqueId(root, this->stub_, pimpl, &nccl_unique_id_);

src/common/cuda_dr_utils.cc

@@ -1,10 +1,11 @@
 /**
- * Copyright 2024-2025, XGBoost contributors
+ * Copyright 2024-2026, XGBoost contributors
  */
 #if defined(XGBOOST_USE_CUDA)
 #include "cuda_dr_utils.h"
 
 #include <algorithm>  // for max
+#include <array>      // for array
 #include <charconv>   // for from_chars
 #include <cstdint>    // for int32_t
 #include <cstring>    // for memset
@@ -210,9 +211,14 @@ namespace detail {
     return -1;
   }
 
+  std::array<unsigned char, curt::kUuidLength> uuid;
+  // Select the current GPU to query.
+  curt::GetUuid(common::Span{uuid.data(), uuid.size()}, curt::CurrentDevice());
+  auto str_uuid = curt::PrintUuid(common::Span{uuid.data(), uuid.size()});
   // See test for example output from smi.
-  auto cmd = "nvidia-smi c2c -s -i 0";  // Select the first GPU to query.
+  auto cmd = "nvidia-smi c2c -s -i " + str_uuid;
   auto out = common::CmdOutput(StringView{cmd});
+  LOG(DEBUG) << "c2c:\n" << out << "\n";
   auto cnt = detail::GetC2cLinkCountFromSmiImpl(out);
   return cnt;
 }

src/common/cuda_rt_utils.cc

@@ -1,15 +1,21 @@
 /**
- * Copyright 2015-2025, XGBoost Contributors
+ * Copyright 2015-2026, XGBoost Contributors
  */
 #include "cuda_rt_utils.h"
 
-#include "cuda_stream.h"  // for StreamRef
+#include <cstring>  // for memcpy
+#include <set>      // for set
+#include <sstream>  // for stringstream
+
+#include "cuda_stream.h"   // for StreamRef
+#include "xgboost/span.h"  // for Span
 
 #if defined(XGBOOST_USE_CUDA)
 #include <cuda_runtime_api.h>
 
 #include <algorithm>  // for max
-#endif                // defined(XGBOOST_USE_CUDA)
+
+#endif  // defined(XGBOOST_USE_CUDA)
 
 #include <cstddef>  // for size_t
 #include <cstdint>  // for int32_t
@@ -114,6 +120,27 @@ void GetDrVersionGlobal(std::int32_t* major, std::int32_t* minor) {
   return !!res;
 }
 
+static_assert(kUuidLength == sizeof(std::declval<cudaDeviceProp>().uuid));
+
+void GetUuid(xgboost::common::Span<unsigned char> uuid, std::int32_t device) {
+  cudaDeviceProp prop{};
+  dh::safe_cuda(cudaGetDeviceProperties(&prop, device));
+  std::memcpy(uuid.data(), static_cast<void*>(&(prop.uuid)), kUuidLength);
+}
+
+[[nodiscard]] std::string PrintUuid(common::Span<unsigned char const, kUuidLength> uuid) {
+  std::set<std::size_t> dash_pos{0, 4, 6, 8, 10};
+  std::stringstream ss;
+  ss << "GPU";
+  for (std::size_t i = 0; i < kUuidLength; ++i) {
+    if (dash_pos.find(i) != dash_pos.cend()) {
+      ss << "-";
+    }
+    ss << std::setw(2) << std::setfill('0') << std::hex << (0xFF & std::uint32_t{uuid[i]});
+  }
+  return ss.str();
+}
+
 void MemcpyAsync(void* dst, const void* src, std::size_t count, StreamRef stream) {
   dh::safe_cuda(cudaMemcpyAsync(dst, src, count, cudaMemcpyDefault, stream));
 }
@@ -154,6 +181,10 @@ void SetDevice(std::int32_t device) {
 
 [[nodiscard]] bool MemoryPoolsSupported(std::int32_t) { return false; }
 
+void GetUuid(xgboost::common::Span<unsigned char>, std::int32_t) { common::AssertGPUSupport(); }
+
+[[nodiscard]] std::string PrintUuid(common::Span<unsigned char const, kUuidLength>) { return {}; }
+
 void MemcpyAsync(void*, const void*, std::size_t, StreamRef) { common::AssertGPUSupport(); }
 
 #endif  // !defined(XGBOOST_USE_CUDA)

src/common/cuda_rt_utils.h

@@ -1,11 +1,13 @@
 /**
- * Copyright 2024-2025, XGBoost contributors
+ * Copyright 2024-2026, XGBoost contributors
  */
 #pragma once
 #include <cstddef>  // for size_t
 #include <cstdint>  // for int32_t
+#include <string>   // for string
 
-#include "cuda_stream.h"  // for StreamRef
+#include "cuda_stream.h"   // for StreamRef
+#include "xgboost/span.h"  // for Span
 
 namespace xgboost::curt {
 std::int32_t AllVisibleGPUs();
@@ -42,6 +44,12 @@ void GetDrVersionGlobal(std::int32_t* major, std::int32_t* minor);
 
 [[nodiscard]] bool MemoryPoolsSupported(std::int32_t device);
 
+inline constexpr std::size_t kUuidLength = 16;
+
+void GetUuid(common::Span<unsigned char> uuid, std::int32_t device);
+
+[[nodiscard]] std::string PrintUuid(common::Span<unsigned char const, kUuidLength> uuid);
+
 // cudaMemcpyAsync
 void MemcpyAsync(void* dst, const void* src, std::size_t count, StreamRef stream);
 }  // namespace xgboost::curt

src/data/batch_utils.cc

@@ -51,16 +51,6 @@ void CheckParam(BatchParam const& init, BatchParam const& param) {
   using xgboost::cuda_impl::CachePageRatio;
 
   auto lc = cudr::GetC2cLinkCountFromSmiGlobal();
-  if (lc >= 10) {
-    // >= 10, life is easy.
-    if (CachePageBytesIsAuto(min_cache_page_bytes)) {
-      min_cache_page_bytes = n_d_bytes * CachePageRatio();
-    }
-    if (HostRatioIsAuto(cache_host_ratio)) {
-      cache_host_ratio = 1.0;
-    }
-    return {cache_host_ratio, min_cache_page_bytes};
-  }
 
   /**
    * Configure the min_cache_page_bytes
@@ -96,7 +86,6 @@ void CheckParam(BatchParam const& init, BatchParam const& param) {
     auto h_cache_nbytes = n_cache_bytes - d_cache_nbytes * 0.85;
     cache_host_ratio = static_cast<double>(h_cache_nbytes) / static_cast<double>(n_cache_bytes);
     if (lc > 0) {
-      // 0 < lc < 10, C2C is available, but with reduced link count.
       // No need to exceed half in practice.
       cache_host_ratio = std::max(cache_host_ratio, 0.5);
     }

tests/cpp/common/test_cuda_rt_utils.cu

@@ -1,15 +1,27 @@
 /**
- * Copyright 2025, XGBoost contributors
+ * Copyright 2025-2026, XGBoost contributors
  */
 
 #include <gtest/gtest.h>
 
+#include <array>    // for array
 #include <cstdint>  // for int32_t
 #include <set>      // for set
 
+#include "../../../src/common/cuda_rt_utils.h"
 #include "../../../src/common/cuda_stream_pool.h"
+#include "xgboost/span.h"  // for Span
 
 namespace xgboost::curt {
+TEST(RtUtils, Uuid) {
+  std::array<unsigned char, kUuidLength> uuid;
+  GetUuid(uuid, 0);
+  auto str = PrintUuid(uuid);
+  ASSERT_EQ(str.substr(0, 4), "GPU-");
+  ASSERT_EQ(str.length(), 40);
+  ASSERT_EQ(str.size(), str.length());
+}
+
 TEST(RtUtils, StreamPool) {
   auto n_streams = 16;
   auto pool = std::make_unique<StreamPool>(n_streams);