Don't estimate sketch batch size when rmm is used. (#6807)

2021-03-31 15:29:56 +08:00
parent 10ae0f9511
commit 3039dd194b
2 changed files with 9 additions and 0 deletions
--- a/src/common/hist_util.cu
+++ b/src/common/hist_util.cu
@@ -93,6 +93,11 @@ size_t SketchBatchNumElements(size_t sketch_batch_num_elements,
                              bst_row_t num_rows, bst_feature_t columns,
                              size_t nnz, int device,
                              size_t num_cuts, bool has_weight) {
+#if defined(XGBOOST_USE_RMM) && XGBOOST_USE_RMM == 1
+  // device available memory is not accurate when rmm is used.
+  return nnz;
+#endif  // defined(XGBOOST_USE_RMM) && XGBOOST_USE_RMM == 1
+
  if (sketch_batch_num_elements == 0) {
    auto required_memory = RequiredMemory(num_rows, columns, nnz, num_cuts, has_weight);
    // use up to 80% of available space