Reduce memory in GPU quantile sketch for external memory (#12047)

Author: RAMitchell

include/xgboost/c_api.h

@@ -529,12 +529,8 @@ XGB_DLL int XGDMatrixCreateFromCallback(DataIterHandle iter, DMatrixHandle proxy
  *   - nthread (optional): Number of threads used for initializing DMatrix.
  *   - max_bin (optional): Maximum number of bins for building histogram. Must be consistent with
  *                         the corresponding booster training parameter.
- *   - max_quantile_blocks (optional): For GPU-based inputs, XGBoost handles incoming
- *       batches with multiple growing substreams. This parameter sets the maximum number
- *       of batches before XGBoost can cut the sub-stream and create a new one. This can
- *       help bound the memory usage. By default, XGBoost grows new sub-streams
- *       exponentially until batches are exhausted. Only used for the training dataset and
- *       the default is None (unbounded).
+ *   - max_quantile_blocks (optional, deprecated): This parameter no longer has any effect and
+ *       will be removed in a future release.
  * @param out      The created Quantile DMatrix.
  *
  * @return 0 when success, -1 when failure happens
@@ -570,12 +566,8 @@ XGB_DLL int XGQuantileDMatrixCreateFromCallback(DataIterHandle iter, DMatrixHand
  *   - min_cache_page_bytes (optional): The minimum number of bytes for each internal GPU
  *      page. Set to 0 to disable page concatenation. Automatic configuration if the
  *      parameter is not provided or set to None.
- *   - max_quantile_blocks (optional): For GPU-based inputs, XGBoost handles incoming
- *       batches with multiple growing substreams. This parameter sets the maximum number
- *       of batches before XGBoost can cut the sub-stream and create a new one. This can
- *       help bound the memory usage. By default, XGBoost grows new sub-streams
- *       exponentially until batches are exhausted. Only used for the training dataset and
- *       the default is None (unbounded).
+ *   - max_quantile_blocks (optional, deprecated): This parameter no longer has any effect and
+ *       will be removed in a future release.
  * - cache_host_ratio (optioinal): For GPU-based inputs, XGBoost can split the cache into
  *      host and device portitions to reduce the data transfer overhead. This parameter
  *      specifies the size of host cache compared to the size of the entire cache:

include/xgboost/data.h

@@ -673,7 +673,7 @@ class DMatrix {
             typename XGDMatrixCallbackNext>
   static DMatrix* Create(DataIterHandle iter, DMatrixHandle proxy, std::shared_ptr<DMatrix> ref,
                          DataIterResetCallback* reset, XGDMatrixCallbackNext* next, float missing,
-                         std::int32_t nthread, bst_bin_t max_bin, std::int64_t max_quantile_blocks);
+                         std::int32_t nthread, bst_bin_t max_bin);
 
   /**
    * @brief Create an external memory DMatrix with callbacks.
@@ -707,8 +707,7 @@ class DMatrix {
             typename XGDMatrixCallbackNext>
   static DMatrix* Create(DataIterHandle iter, DMatrixHandle proxy, std::shared_ptr<DMatrix> ref,
                          DataIterResetCallback* reset, XGDMatrixCallbackNext* next,
-                         bst_bin_t max_bin, std::int64_t max_quantile_blocks,
-                         ExtMemConfig const& config);
+                         bst_bin_t max_bin, ExtMemConfig const& config);
 
   virtual DMatrix* Slice(common::Span<int32_t const> ridxs) = 0;
 

python-package/xgboost/core.py

@@ -1465,19 +1465,12 @@ class QuantileDMatrix(DMatrix, _RefMixIn):
         applied to the validation/test data
 
     max_quantile_batches :
-        For GPU-based inputs from an iterator, XGBoost handles incoming batches with
-        multiple growing sub-streams. This parameter sets the maximum number of batches
-        before XGBoost can cut a sub-stream and create a new one. This can help bound
-        the memory usage. By default, XGBoost grows a sub-stream exponentially until
-        batches are exhausted. This option is only used for the training dataset and the
-        default is None (unbounded). Lastly, if the `data` is a single batch instead of
-        an iterator, this parameter has no effect.
+        Deprecated. This parameter no longer has any effect and will be removed in a
+        future release.
 
         .. versionadded:: 3.0.0
 
-        .. warning::
-
-            This is an experimental parameter and subject to change.
+        .. deprecated:: 3.3.0
 
     """
 
@@ -1642,7 +1635,7 @@ def __init__(  # pylint: disable=super-init-not-called
             A user-defined :py:class:`DataIter` for loading data.
 
         max_quantile_batches :
-            See :py:class:`QuantileDMatrix`.
+            Deprecated. See :py:class:`QuantileDMatrix`.
 
         cache_host_ratio :
 

src/c_api/c_api.cc

@@ -367,6 +367,15 @@ std::shared_ptr<DMatrix> GetRefDMatrix(DataIterHandle ref) {
   }
   return _ref;
 }
+
+void WarnDeprecatedMaxQuantileBlocks(Json const &config) {
+  auto const &obj = get<Object const>(config);
+  auto it = obj.find("max_quantile_blocks");
+  if (it != obj.cend() && !IsA<Null>(it->second)) {
+    LOG(WARNING) << "`max_quantile_blocks` is deprecated and has no effect. "
+                    "The parameter will be removed in a future release.";
+  }
+}
 }  // namespace
 
 XGB_DLL int XGQuantileDMatrixCreateFromCallback(DataIterHandle iter, DMatrixHandle proxy,
@@ -378,18 +387,17 @@ XGB_DLL int XGQuantileDMatrixCreateFromCallback(DataIterHandle iter, DMatrixHand
 
   xgboost_CHECK_C_ARG_PTR(config);
   auto jconfig = Json::Load(StringView{config});
+  WarnDeprecatedMaxQuantileBlocks(jconfig);
   auto missing = GetMissing(jconfig);
   auto n_threads = OptionalArg<Integer, int64_t>(jconfig, "nthread", 0);
   auto max_bin = OptionalArg<Integer, int64_t>(jconfig, "max_bin", 256);
-  auto max_quantile_blocks = OptionalArg<Integer, std::int64_t>(
-      jconfig, "max_quantile_blocks", std::numeric_limits<std::int64_t>::max());
 
   xgboost_CHECK_C_ARG_PTR(next);
   xgboost_CHECK_C_ARG_PTR(reset);
   xgboost_CHECK_C_ARG_PTR(out);
 
-  *out = new std::shared_ptr<xgboost::DMatrix>{xgboost::DMatrix::Create(
-      iter, proxy, p_ref, reset, next, missing, n_threads, max_bin, max_quantile_blocks)};
+  *out = new std::shared_ptr<xgboost::DMatrix>{
+      xgboost::DMatrix::Create(iter, proxy, p_ref, reset, next, missing, n_threads, max_bin)};
   API_END();
 }
 
@@ -403,15 +411,14 @@ XGB_DLL int XGExtMemQuantileDMatrixCreateFromCallback(DataIterHandle iter, DMatr
 
   xgboost_CHECK_C_ARG_PTR(config);
   auto jconfig = Json::Load(StringView{config});
+  WarnDeprecatedMaxQuantileBlocks(jconfig);
   auto missing = GetMissing(jconfig);
   std::int32_t n_threads = OptionalArg<Integer, std::int64_t>(jconfig, "nthread", 0);
   auto max_bin = OptionalArg<Integer, std::int64_t>(jconfig, "max_bin", 256);
   auto on_host = OptionalArg<Boolean>(jconfig, "on_host", false);
   std::string cache = RequiredArg<String>(jconfig, "cache_prefix", __func__);
   auto min_cache_page_bytes = OptionalArg<Integer, std::int64_t>(jconfig, "min_cache_page_bytes",
                                                                  cuda_impl::AutoCachePageBytes());
-  auto max_quantile_blocks = OptionalArg<Integer, std::int64_t>(
-      jconfig, "max_quantile_blocks", std::numeric_limits<std::int64_t>::max());
   auto cache_host_ratio =
       OptionalArg<Number, float>(jconfig, "cache_host_ratio", cuda_impl::AutoHostRatio());
 
@@ -421,8 +428,8 @@ XGB_DLL int XGExtMemQuantileDMatrixCreateFromCallback(DataIterHandle iter, DMatr
 
   auto config =
       ExtMemConfig{cache, on_host, cache_host_ratio, min_cache_page_bytes, missing, n_threads};
-  *out = new std::shared_ptr<xgboost::DMatrix>{xgboost::DMatrix::Create(
-      iter, proxy, p_ref, reset, next, max_bin, max_quantile_blocks, config)};
+  *out = new std::shared_ptr<xgboost::DMatrix>{
+      xgboost::DMatrix::Create(iter, proxy, p_ref, reset, next, max_bin, config)};
   API_END();
 }
 

src/common/hist_util.cu

@@ -367,8 +367,7 @@ HistogramCuts DeviceSketchWithHessian(Context const* ctx, DMatrix* p_fmat, bst_b
   auto d_weight = UnifyWeight(cuctx, info, hessian, &weight);
 
   HistogramCuts cuts;
-  SketchContainer sketch_container(info.feature_types, max_bin, info.num_col_, info.num_row_,
-                                   ctx->Device());
+  SketchContainer sketch_container(info.feature_types, max_bin, info.num_col_, ctx->Device());
   CHECK_EQ(has_weight || !hessian.empty(), !d_weight.empty());
   for (const auto& page : p_fmat->GetBatches<SparsePage>()) {
     std::size_t page_nnz = page.data.Size();

src/common/quantile.cu

@@ -505,18 +505,14 @@ void SketchContainer::AllReduce(Context const *ctx, bool is_column_split) {
   }
 
   timer_.Start(__func__);
-  // Reduce the overhead on syncing.
-  bst_idx_t global_sum_rows = num_rows_;
-  auto rc = collective::Allreduce(ctx, linalg::MakeVec(&global_sum_rows, 1), collective::Op::kSum);
-  SafeColl(rc);
-  bst_idx_t intermediate_num_cuts =
-      std::min(global_sum_rows, static_cast<size_t>(num_bins_ * kFactor));
+  // Bound local sketch size before exchanging data across workers.
+  auto intermediate_num_cuts = static_cast<bst_idx_t>(num_bins_ * kFactor);
   this->Prune(ctx, intermediate_num_cuts);
 
   auto d_columns_ptr = this->columns_ptr_.ConstDeviceSpan();
   CHECK_EQ(d_columns_ptr.size(), num_columns_ + 1);
   size_t n = d_columns_ptr.size();
-  rc = collective::Allreduce(ctx, linalg::MakeVec(&n, 1), collective::Op::kMax);
+  auto rc = collective::Allreduce(ctx, linalg::MakeVec(&n, 1), collective::Op::kMax);
   SafeColl(rc);
   CHECK_EQ(n, d_columns_ptr.size()) << "Number of columns differs across workers";
 
@@ -560,8 +556,7 @@ void SketchContainer::AllReduce(Context const *ctx, bool is_column_split) {
   timer_.Stop(__func__);
 
   // Merge them into a new sketch.
-  SketchContainer new_sketch(this->feature_types_, num_bins_, this->num_columns_, global_sum_rows,
-                             ctx->Device());
+  SketchContainer new_sketch(this->feature_types_, num_bins_, this->num_columns_, ctx->Device());
   for (size_t i = 0; i < allworkers.size(); ++i) {
     auto worker = allworkers[i];
     auto worker_ptr =

src/common/quantile.cuh

@@ -46,7 +46,6 @@ class SketchContainer {
  private:
   Monitor timer_;
   HostDeviceVector<FeatureType> feature_types_;
-  bst_idx_t num_rows_;
   bst_feature_t num_columns_;
   int32_t num_bins_;
 
@@ -97,12 +96,11 @@ class SketchContainer {
    *
    * \param max_bin     Maximum number of bins per columns
    * \param num_columns Total number of columns in dataset.
-   * \param num_rows    Total number of rows in known dataset (typically the rows in current worker).
    * \param device      GPU ID.
    */
   SketchContainer(HostDeviceVector<FeatureType> const& feature_types, bst_bin_t max_bin,
-                  bst_feature_t num_columns, bst_idx_t num_rows, DeviceOrd device)
-      : num_rows_{num_rows}, num_columns_{num_columns}, num_bins_{max_bin} {
+                  bst_feature_t num_columns, DeviceOrd device)
+      : num_columns_{num_columns}, num_bins_{max_bin} {
     CHECK(device.IsCUDA());
     // Initialize Sketches for this dmatrix
     this->columns_ptr_.SetDevice(device);

src/data/data.cc

@@ -974,9 +974,8 @@ template <typename DataIterHandle, typename DMatrixHandle, typename DataIterRese
           typename XGDMatrixCallbackNext>
 DMatrix* DMatrix::Create(DataIterHandle iter, DMatrixHandle proxy, std::shared_ptr<DMatrix> ref,
                          DataIterResetCallback* reset, XGDMatrixCallbackNext* next, float missing,
-                         int nthread, bst_bin_t max_bin, std::int64_t max_quantile_blocks) {
-  return new data::IterativeDMatrix(iter, proxy, ref, reset, next, missing, nthread, max_bin,
-                                    max_quantile_blocks);
+                         int nthread, bst_bin_t max_bin) {
+  return new data::IterativeDMatrix(iter, proxy, ref, reset, next, missing, nthread, max_bin);
 }
 
 template <typename DataIterHandle, typename DMatrixHandle, typename DataIterResetCallback,
@@ -990,17 +989,16 @@ template <typename DataIterHandle, typename DMatrixHandle, typename DataIterRese
           typename XGDMatrixCallbackNext>
 DMatrix* DMatrix::Create(DataIterHandle iter, DMatrixHandle proxy, std::shared_ptr<DMatrix> ref,
                          DataIterResetCallback* reset, XGDMatrixCallbackNext* next,
-                         bst_bin_t max_bin, std::int64_t max_quantile_blocks,
-                         ExtMemConfig const& config) {
-  return new data::ExtMemQuantileDMatrix{
-      iter, proxy, ref, reset, next, max_bin, max_quantile_blocks, config};
+                         bst_bin_t max_bin, ExtMemConfig const& config) {
+  return new data::ExtMemQuantileDMatrix{iter, proxy, ref, reset, next, max_bin, config};
 }
 
-template DMatrix*
-DMatrix::Create<DataIterHandle, DMatrixHandle, DataIterResetCallback, XGDMatrixCallbackNext>(
-    DataIterHandle iter, DMatrixHandle proxy, std::shared_ptr<DMatrix> ref,
-    DataIterResetCallback* reset, XGDMatrixCallbackNext* next, float missing, int nthread,
-    int max_bin, std::int64_t max_quantile_blocks);
+template DMatrix* DMatrix::Create<DataIterHandle, DMatrixHandle, DataIterResetCallback,
+                                  XGDMatrixCallbackNext>(DataIterHandle iter, DMatrixHandle proxy,
+                                                         std::shared_ptr<DMatrix> ref,
+                                                         DataIterResetCallback* reset,
+                                                         XGDMatrixCallbackNext* next, float missing,
+                                                         int nthread, int max_bin);
 
 template DMatrix* DMatrix::Create<DataIterHandle, DMatrixHandle, DataIterResetCallback,
                                   XGDMatrixCallbackNext>(DataIterHandle iter, DMatrixHandle proxy,
@@ -1011,7 +1009,7 @@ template DMatrix* DMatrix::Create<DataIterHandle, DMatrixHandle, DataIterResetCa
 template DMatrix*
 DMatrix::Create<DataIterHandle, DMatrixHandle, DataIterResetCallback, XGDMatrixCallbackNext>(
     DataIterHandle, DMatrixHandle, std::shared_ptr<DMatrix>, DataIterResetCallback*,
-    XGDMatrixCallbackNext*, bst_bin_t, std::int64_t, ExtMemConfig const&);
+    XGDMatrixCallbackNext*, bst_bin_t, ExtMemConfig const&);
 
 template <typename AdapterT>
 DMatrix* DMatrix::Create(AdapterT* adapter, float missing, int nthread, const std::string&,

src/data/extmem_quantile_dmatrix.cc

@@ -24,7 +24,6 @@ ExtMemQuantileDMatrix::ExtMemQuantileDMatrix(DataIterHandle iter_handle, DMatrix
                                              std::shared_ptr<DMatrix> ref,
                                              DataIterResetCallback *reset,
                                              XGDMatrixCallbackNext *next, bst_bin_t max_bin,
-                                             std::int64_t max_quantile_blocks,
                                              ExtMemConfig const &config)
     : cache_prefix_{config.cache}, on_host_{config.on_host} {
   cache_prefix_ = MakeCachePrefix(cache_prefix_);
@@ -45,7 +44,7 @@ ExtMemQuantileDMatrix::ExtMemQuantileDMatrix(DataIterHandle iter_handle, DMatrix
     this->InitFromCPU(&ctx, iter, proxy, p, config.missing, ref);
   } else {
     p.n_prefetch_batches = ::xgboost::cuda_impl::DftPrefetchBatches();
-    this->InitFromCUDA(&ctx, iter, proxy, p, ref, max_quantile_blocks, config);
+    this->InitFromCUDA(&ctx, iter, proxy, p, ref, config);
   }
   this->batch_ = p;
   this->fmat_ctx_ = ctx;
@@ -144,8 +143,7 @@ BatchSet<GHistIndexMatrix> ExtMemQuantileDMatrix::GetGradientIndex(Context const
 #if !defined(XGBOOST_USE_CUDA)
 void ExtMemQuantileDMatrix::InitFromCUDA(
     Context const *, std::shared_ptr<DataIterProxy<DataIterResetCallback, XGDMatrixCallbackNext>>,
-    DMatrixHandle, BatchParam const &, std::shared_ptr<DMatrix>, std::int64_t,
-    ExtMemConfig const &) {
+    DMatrixHandle, BatchParam const &, std::shared_ptr<DMatrix>, ExtMemConfig const &) {
   common::AssertGPUSupport();
 }
 

src/data/extmem_quantile_dmatrix.cu

@@ -41,7 +41,7 @@ void ExtMemQuantileDMatrix::InitFromCUDA(
     Context const *ctx,
     std::shared_ptr<DataIterProxy<DataIterResetCallback, XGDMatrixCallbackNext>> iter,
     DMatrixHandle proxy_handle, BatchParam const &p, std::shared_ptr<DMatrix> ref,
-    std::int64_t max_quantile_blocks, ExtMemConfig const &config) {
+    ExtMemConfig const &config) {
   xgboost_NVTX_FN_RANGE();
 
   // A handle passed to external iterator.
@@ -54,7 +54,7 @@ void ExtMemQuantileDMatrix::InitFromCUDA(
   auto cuts = std::make_shared<common::HistogramCuts>();
   ExternalDataInfo ext_info;
   cuda_impl::MakeSketches(ctx, iter.get(), proxy, ref, p, config.missing, cuts, this->info_,
-                          max_quantile_blocks, &ext_info);
+                          &ext_info);
   ext_info.SetInfo(ctx, true, &this->info_);
 
   /**