From 507a7635531db600a28dd06ab4480269abd64fb8 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Tue, 11 Aug 2026 01:53:28 +0000 Subject: [PATCH 1/2] cuml#8394 squashed --- .../batched-levelalgo/builder.cuh | 89 ++- .../kernels/builder_kernels.cuh | 25 +- .../kernels/builder_kernels_impl.cuh | 65 +- .../kernels/classification-double.cu | 10 +- .../kernels/classification-float.cu | 10 +- .../kernels/regression-double.cu | 10 +- .../kernels/regression-float.cu | 10 +- .../kernels/weighted-classification-double.cu | 10 +- .../kernels/weighted-classification-float.cu | 10 +- .../kernels/weighted-regression-double.cu | 10 +- .../kernels/weighted-regression-float.cu | 10 +- .../batched-levelalgo/quantiles.cuh | 14 +- cpp/src/randomforest/randomforest.cuh | 28 +- cpp/tests/CMakeLists.txt | 1 + cpp/tests/mg/rf_test.cu | 681 ++++++++++++++++++ cpp/tests/sg/rf_test.cu | 15 + 16 files changed, 912 insertions(+), 86 deletions(-) create mode 100644 cpp/tests/mg/rf_test.cu diff --git a/cpp/src/decisiontree/batched-levelalgo/builder.cuh b/cpp/src/decisiontree/batched-levelalgo/builder.cuh index 3bad9563fa..68bd00b24e 100644 --- a/cpp/src/decisiontree/batched-levelalgo/builder.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/builder.cuh @@ -47,16 +47,20 @@ class NodeQueue { std::deque work_items_; public: - NodeQueue(DecisionTreeParams params, size_t max_nodes, size_t sampled_rows, int num_outputs) + NodeQueue(DecisionTreeParams params, + size_t max_nodes, + size_t local_sampled_rows, + std::int64_t global_sampled_rows, + int num_outputs) : params(params), tree(std::make_shared>()) { tree->num_outputs = num_outputs; tree->sparsetree.reserve(max_nodes); - tree->sparsetree.emplace_back(NodeT::CreateLeafNode(sampled_rows)); + tree->sparsetree.emplace_back(NodeT::CreateLeafNode(global_sampled_rows)); tree->leaf_counter = 1; tree->depth_counter = 0; node_instances_.reserve(max_nodes); - node_instances_.emplace_back(InstanceRange{0, sampled_rows}); + node_instances_.emplace_back(InstanceRange{0, local_sampled_rows}); if (this->IsExpandable(tree->sparsetree.back(), 0)) { work_items_.emplace_back(NodeWorkItem{0, 0, node_instances_.back()}); } @@ -100,15 +104,15 @@ class NodeQueue { if (params.max_leaves != -1 && tree->leaf_counter >= params.max_leaves) break; using NodeCountT = decltype(std::declval().InstanceCount()); + auto const parent_count = tree->sparsetree.at(item.idx).InstanceCount(); auto const local_left_count = ML::narrow_cast(split.local_nLeft); // parent - tree->sparsetree.at(item.idx) = - NodeT::CreateSplitNode(split.colid, - split.quesval, - split.best_metric_val, - int64_t(tree->sparsetree.size()), - ML::narrow_cast(parent_range.count)); + tree->sparsetree.at(item.idx) = NodeT::CreateSplitNode(split.colid, + split.quesval, + split.best_metric_val, + int64_t(tree->sparsetree.size()), + parent_count); tree->leaf_counter++; // left tree->sparsetree.emplace_back( @@ -122,9 +126,8 @@ class NodeQueue { } // right - tree->sparsetree.emplace_back(NodeT::CreateLeafNode( - ML::checked_sub(tree->sparsetree.at(item.idx).InstanceCount(), - ML::narrow_cast(split.global_nLeft)))); + tree->sparsetree.emplace_back(NodeT::CreateLeafNode(ML::checked_sub( + parent_count, ML::narrow_cast(split.global_nLeft)))); node_instances_.emplace_back( InstanceRange{ML::checked_add(parent_range.begin, local_left_count), ML::checked_sub(parent_range.count, local_left_count)}); @@ -391,8 +394,11 @@ struct Builder { { raft::common::nvtx::range fun_scope("Builder::train @builder.cuh [batched-levelalgo]"); MLCommon::TimerCPU timer; - NodeQueue queue( - params, this->maxNodes(), dataset.n_sampled_rows, dataset.num_outputs); + NodeQueue queue(params, + this->maxNodes(), + dataset.n_sampled_rows, + this->globalSampledRows(), + dataset.num_outputs); while (queue.HasWork()) { auto work_items = queue.Pop(); auto [splits_host_ptr, splits_count] = doSplit(work_items); @@ -425,6 +431,22 @@ struct Builder { return n_blocks_dimx; } + std::int64_t globalSampledRows() + { + auto global_sampled_rows = static_cast(dataset.n_sampled_rows); + if (!distributed) { return global_sampled_rows; } + + rmm::device_uvector d_sampled_rows(1, builder_stream); + raft::update_device(d_sampled_rows.data(), &global_sampled_rows, 1, builder_stream); + handle.get_comms().allreduce( + d_sampled_rows.data(), d_sampled_rows.data(), 1, raft::comms::op_t::SUM, builder_stream); + ASSERT(handle.get_comms().sync_stream(builder_stream) == raft::comms::status_t::SUCCESS, + "An error occurred in the distributed RF sampled-row-count all-reduce."); + raft::update_host(&global_sampled_rows, d_sampled_rows.data(), 1, builder_stream); + handle.sync_stream(builder_stream); + return global_sampled_rows; + } + auto doSplit(const std::vector& work_items) { raft::common::nvtx::range fun_scope("Builder::doSplit @builder.cuh [batched-levelalgo]"); @@ -659,9 +681,13 @@ struct Builder { tree->vector_leaf.resize(vector_leaf_size); ASSERT(tree->sparsetree.size() == instance_ranges.size(), "Expected instance range for each node"); - // do this in batch to reduce peak memory usage in extreme cases - std::size_t max_batch_size = min(std::size_t{100000}, tree->sparsetree.size()); - auto max_leaf_values = ML::checked_mul(max_batch_size, dataset.num_outputs); + // Reuse the split histogram and packed reduction workspaces for leaf statistics. Cap the + // number of nodes so each leaf batch fits in those workspaces. + auto max_leaf_nodes_in_workspace = + ML::checked_mul(params.max_batch_size, params.max_n_bins, n_blks_for_cols); + std::size_t max_batch_size = + std::min(std::size_t{100000}, std::min(tree->sparsetree.size(), max_leaf_nodes_in_workspace)); + auto max_leaf_values = ML::checked_mul(max_batch_size, dataset.num_outputs); rmm::device_uvector d_tree(max_batch_size, builder_stream); rmm::device_uvector d_instance_ranges(max_batch_size, builder_stream); rmm::device_uvector d_leaves(max_leaf_values, builder_stream); @@ -675,17 +701,26 @@ struct Builder { raft::update_device( d_instance_ranges.data(), instance_ranges.data() + batch_begin, batch_size, builder_stream); - auto leaves_bytes = ML::checked_mul(sizeof(DataT), d_leaves.size()); - RAFT_CUDA_TRY(cudaMemsetAsync(d_leaves.data(), 0, leaves_bytes, builder_stream)); + auto leaf_histogram_count = ML::checked_mul(batch_size, dataset.num_outputs); + auto leaf_histogram_bytes = ML::checked_mul(sizeof(BinT), leaf_histogram_count); + auto leaf_batch_size = ML::narrow_cast(batch_size); + RAFT_CUDA_TRY(cudaMemsetAsync(histograms, 0, leaf_histogram_bytes, builder_stream)); size_t smem_size = ML::checked_mul(sizeof(BinT), dataset.num_outputs); - launchLeafKernel(objective, - dataset, - d_tree.data(), - d_instance_ranges.data(), - d_leaves.data(), - batch_size, - smem_size, - builder_stream); + launchBuildLeafHistogramsKernel(objective, + dataset, + d_tree.data(), + d_instance_ranges.data(), + histograms, + leaf_batch_size, + smem_size, + builder_stream); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + if (distributed) { allReduceHistograms(histograms, leaf_histogram_count); } + + launchFinalizeLeafKernel( + histograms, d_leaves.data(), dataset.num_outputs, leaf_batch_size, builder_stream); + RAFT_CUDA_TRY(cudaPeekAtLastError()); auto leaf_offset = ML::checked_mul(batch_begin, dataset.num_outputs); auto leaf_count = ML::checked_mul(batch_size, dataset.num_outputs); raft::update_host( diff --git a/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh b/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh index 68635df491..63135f4c24 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh @@ -107,15 +107,22 @@ void launchNodeSplitKernel(const Dataset& dataset, std::int64_t* partition_row_ids, cudaStream_t builder_stream); -template -void launchLeafKernel(ObjectiveT objective, - DatasetT& dataset, - const NodeT* tree, - const InstanceRange* instance_ranges, - DataT* leaves, - int batch_size, - size_t smem_size, - cudaStream_t builder_stream); +template +void launchBuildLeafHistogramsKernel(ObjectiveT objective, + DatasetT& dataset, + const NodeT* tree, + const InstanceRange* instance_ranges, + typename ObjectiveT::BinT* leaf_histograms, + int batch_size, + size_t smem_size, + cudaStream_t builder_stream); + +template +void launchFinalizeLeafKernel(const typename ObjectiveT::BinT* leaf_histograms, + DataT* leaves, + int num_outputs, + int batch_size, + cudaStream_t builder_stream); template void launchBuildHistogramsKernel(typename ObjectiveT::BinT* histograms, std::int64_t n_bins, diff --git a/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels_impl.cuh b/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels_impl.cuh index 9929f8e73c..b4a788e44b 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels_impl.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels_impl.cuh @@ -212,12 +212,12 @@ void launchNodeSplitKernel(const Dataset& dataset, dataset, work_items, splits, workload_info, partition_row_ids); } -template -static __global__ void leafKernel(ObjectiveT objective, - DatasetT dataset, - const NodeT* tree, - const InstanceRange* instance_ranges, - DataT* leaves) +template +static __global__ void buildLeafHistogramsKernel(ObjectiveT objective, + DatasetT dataset, + const NodeT* tree, + const InstanceRange* instance_ranges, + typename ObjectiveT::BinT* leaf_histograms) { using BinT = typename ObjectiveT::BinT; extern __shared__ char shared_memory[]; @@ -237,25 +237,48 @@ static __global__ void leafKernel(ObjectiveT objective, objective.IncrementHistogram(histogram, 1, 0, label, dataset, row); } __syncthreads(); - if (tid == 0) { - ObjectiveT::SetLeafVector( - histogram, dataset.num_outputs, leaves + dataset.num_outputs * node_id); + for (int i = tid; i < dataset.num_outputs; i += blockDim.x) { + leaf_histograms[dataset.num_outputs * node_id + i] = histogram[i]; } } -template -void launchLeafKernel(ObjectiveT objective, - DatasetT& dataset, - const NodeT* tree, - const InstanceRange* instance_ranges, - DataT* leaves, - int batch_size, - size_t smem_size, - cudaStream_t builder_stream) +template +void launchBuildLeafHistogramsKernel(ObjectiveT objective, + DatasetT& dataset, + const NodeT* tree, + const InstanceRange* instance_ranges, + typename ObjectiveT::BinT* leaf_histograms, + int batch_size, + size_t smem_size, + cudaStream_t builder_stream) +{ + auto num_blocks = ML::narrow_cast(batch_size); + buildLeafHistogramsKernel<<>>( + objective, dataset, tree, instance_ranges, leaf_histograms); +} + +template +static __global__ void finalizeLeafKernel(const typename ObjectiveT::BinT* leaf_histograms, + DataT* leaves, + int num_outputs, + int batch_size) +{ + auto node_id = int(blockIdx.x) * blockDim.x + threadIdx.x; + if (node_id >= batch_size) return; + ObjectiveT::SetLeafVector( + leaf_histograms + num_outputs * node_id, num_outputs, leaves + num_outputs * node_id); +} + +template +void launchFinalizeLeafKernel(const typename ObjectiveT::BinT* leaf_histograms, + DataT* leaves, + int num_outputs, + int batch_size, + cudaStream_t builder_stream) { - int num_blocks = batch_size; - leafKernel<<>>( - objective, dataset, tree, instance_ranges, leaves); + auto num_blocks = ML::narrow_cast(raft::ceildiv(batch_size, TPB_DEFAULT)); + finalizeLeafKernel<<>>( + leaf_histograms, leaves, num_outputs, batch_size); } /** diff --git a/cpp/src/decisiontree/batched-levelalgo/kernels/classification-double.cu b/cpp/src/decisiontree/batched-levelalgo/kernels/classification-double.cu index 233fc6f9ff..dfcca11901 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/classification-double.cu +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/classification-double.cu @@ -17,16 +17,22 @@ using DatasetT = Dataset; using NodeT = SparseTreeNode; // Explicit instantiations are split across separate .cu files to increase compilation parallelism. -template void launchLeafKernel( +template void launchBuildLeafHistogramsKernel( ObjectiveT objective, DatasetT& dataset, const NodeT* tree, const InstanceRange* instance_ranges, - DataT* leaves, + BinT* leaf_histograms, int batch_size, size_t smem_size, cudaStream_t builder_stream); +template void launchFinalizeLeafKernel(const BinT* leaf_histograms, + DataT* leaves, + int num_outputs, + int batch_size, + cudaStream_t builder_stream); + // Explicit instantiations are split across separate .cu files to increase compilation parallelism. template void launchBuildHistogramsKernel( BinT* histograms, diff --git a/cpp/src/decisiontree/batched-levelalgo/kernels/classification-float.cu b/cpp/src/decisiontree/batched-levelalgo/kernels/classification-float.cu index 7b75c462fb..17fe5d719d 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/classification-float.cu +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/classification-float.cu @@ -17,16 +17,22 @@ using DatasetT = Dataset; using NodeT = SparseTreeNode; // Explicit instantiations are split across separate .cu files to increase compilation parallelism. -template void launchLeafKernel( +template void launchBuildLeafHistogramsKernel( ObjectiveT objective, DatasetT& dataset, const NodeT* tree, const InstanceRange* instance_ranges, - DataT* leaves, + BinT* leaf_histograms, int batch_size, size_t smem_size, cudaStream_t builder_stream); +template void launchFinalizeLeafKernel(const BinT* leaf_histograms, + DataT* leaves, + int num_outputs, + int batch_size, + cudaStream_t builder_stream); + // Explicit instantiations are split across separate .cu files to increase compilation parallelism. template void launchBuildHistogramsKernel( BinT* histograms, diff --git a/cpp/src/decisiontree/batched-levelalgo/kernels/regression-double.cu b/cpp/src/decisiontree/batched-levelalgo/kernels/regression-double.cu index b1e7f54969..b8da442827 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/regression-double.cu +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/regression-double.cu @@ -17,16 +17,22 @@ using DatasetT = Dataset; using NodeT = SparseTreeNode; // Explicit instantiations are split across separate .cu files to increase compilation parallelism. -template void launchLeafKernel( +template void launchBuildLeafHistogramsKernel( ObjectiveT objective, DatasetT& dataset, const NodeT* tree, const InstanceRange* instance_ranges, - DataT* leaves, + BinT* leaf_histograms, int batch_size, size_t smem_size, cudaStream_t builder_stream); +template void launchFinalizeLeafKernel(const BinT* leaf_histograms, + DataT* leaves, + int num_outputs, + int batch_size, + cudaStream_t builder_stream); + // Explicit instantiations are split across separate .cu files to increase compilation parallelism. template void launchBuildHistogramsKernel( BinT* histograms, diff --git a/cpp/src/decisiontree/batched-levelalgo/kernels/regression-float.cu b/cpp/src/decisiontree/batched-levelalgo/kernels/regression-float.cu index 672066cc9d..9d15827518 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/regression-float.cu +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/regression-float.cu @@ -17,16 +17,22 @@ using DatasetT = Dataset; using NodeT = SparseTreeNode; // Explicit instantiations are split across separate .cu files to increase compilation parallelism. -template void launchLeafKernel( +template void launchBuildLeafHistogramsKernel( ObjectiveT objective, DatasetT& dataset, const NodeT* tree, const InstanceRange* instance_ranges, - DataT* leaves, + BinT* leaf_histograms, int batch_size, size_t smem_size, cudaStream_t builder_stream); +template void launchFinalizeLeafKernel(const BinT* leaf_histograms, + DataT* leaves, + int num_outputs, + int batch_size, + cudaStream_t builder_stream); + // Explicit instantiations are split across separate .cu files to increase compilation parallelism. template void launchBuildHistogramsKernel( BinT* histograms, diff --git a/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-classification-double.cu b/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-classification-double.cu index 789a5f21ab..231b60035b 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-classification-double.cu +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-classification-double.cu @@ -17,16 +17,22 @@ using DatasetT = Dataset; using NodeT = SparseTreeNode; // Explicit instantiations are split across separate .cu files to increase compilation parallelism. -template void launchLeafKernel( +template void launchBuildLeafHistogramsKernel( ObjectiveT objective, DatasetT& dataset, const NodeT* tree, const InstanceRange* instance_ranges, - DataT* leaves, + BinT* leaf_histograms, int batch_size, size_t smem_size, cudaStream_t builder_stream); +template void launchFinalizeLeafKernel(const BinT* leaf_histograms, + DataT* leaves, + int num_outputs, + int batch_size, + cudaStream_t builder_stream); + // Explicit instantiations are split across separate .cu files to increase compilation parallelism. template void launchBuildHistogramsKernel( BinT* histograms, diff --git a/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-classification-float.cu b/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-classification-float.cu index 8ebe17eeaf..70f7608fb7 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-classification-float.cu +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-classification-float.cu @@ -17,16 +17,22 @@ using DatasetT = Dataset; using NodeT = SparseTreeNode; // Explicit instantiations are split across separate .cu files to increase compilation parallelism. -template void launchLeafKernel( +template void launchBuildLeafHistogramsKernel( ObjectiveT objective, DatasetT& dataset, const NodeT* tree, const InstanceRange* instance_ranges, - DataT* leaves, + BinT* leaf_histograms, int batch_size, size_t smem_size, cudaStream_t builder_stream); +template void launchFinalizeLeafKernel(const BinT* leaf_histograms, + DataT* leaves, + int num_outputs, + int batch_size, + cudaStream_t builder_stream); + // Explicit instantiations are split across separate .cu files to increase compilation parallelism. template void launchBuildHistogramsKernel( BinT* histograms, diff --git a/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-regression-double.cu b/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-regression-double.cu index d459c46236..b4aa9b7857 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-regression-double.cu +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-regression-double.cu @@ -17,16 +17,22 @@ using DatasetT = Dataset; using NodeT = SparseTreeNode; // Explicit instantiations are split across separate .cu files to increase compilation parallelism. -template void launchLeafKernel( +template void launchBuildLeafHistogramsKernel( ObjectiveT objective, DatasetT& dataset, const NodeT* tree, const InstanceRange* instance_ranges, - DataT* leaves, + BinT* leaf_histograms, int batch_size, size_t smem_size, cudaStream_t builder_stream); +template void launchFinalizeLeafKernel(const BinT* leaf_histograms, + DataT* leaves, + int num_outputs, + int batch_size, + cudaStream_t builder_stream); + // Explicit instantiations are split across separate .cu files to increase compilation parallelism. template void launchBuildHistogramsKernel( BinT* histograms, diff --git a/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-regression-float.cu b/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-regression-float.cu index 21ef9a0c07..78878fb6ea 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-regression-float.cu +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/weighted-regression-float.cu @@ -17,16 +17,22 @@ using DatasetT = Dataset; using NodeT = SparseTreeNode; // Explicit instantiations are split across separate .cu files to increase compilation parallelism. -template void launchLeafKernel( +template void launchBuildLeafHistogramsKernel( ObjectiveT objective, DatasetT& dataset, const NodeT* tree, const InstanceRange* instance_ranges, - DataT* leaves, + BinT* leaf_histograms, int batch_size, size_t smem_size, cudaStream_t builder_stream); +template void launchFinalizeLeafKernel(const BinT* leaf_histograms, + DataT* leaves, + int num_outputs, + int batch_size, + cudaStream_t builder_stream); + // Explicit instantiations are split across separate .cu files to increase compilation parallelism. template void launchBuildHistogramsKernel( BinT* histograms, diff --git a/cpp/src/decisiontree/batched-levelalgo/quantiles.cuh b/cpp/src/decisiontree/batched-levelalgo/quantiles.cuh index f28aefbf18..2af2d25a08 100644 --- a/cpp/src/decisiontree/batched-levelalgo/quantiles.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/quantiles.cuh @@ -154,16 +154,16 @@ CUML_EXPORT QuantileResult computeQuantiles(const raft::handle_t& handle, bool row_major = false) { raft::common::nvtx::push_range("computeQuantiles"); - RAFT_EXPECTS(data != nullptr, "data pointer must not be null"); + auto stream = handle.get_stream(); + bool distributed = raft::resource::comms_initialized(handle) && handle.get_comms().get_size() > 1; + RAFT_EXPECTS(max_n_bins > 0, "max_n_bins must be positive"); - RAFT_EXPECTS(n_rows > 0, "n_rows must be positive"); + RAFT_EXPECTS(distributed ? n_rows >= 0 : n_rows > 0, "n_rows must be positive"); + RAFT_EXPECTS(n_rows == 0 || data != nullptr, "data pointer must not be null"); RAFT_EXPECTS(n_cols > 0, "n_cols must be positive"); RAFT_EXPECTS(oversampling_factor > 0, "oversampling_factor must be positive"); - - auto stream = handle.get_stream(); - bool distributed = raft::resource::comms_initialized(handle) && handle.get_comms().get_size() > 1; - int rank = distributed ? handle.get_comms().get_rank() : 0; - int comm_size = distributed ? handle.get_comms().get_size() : 1; + int rank = distributed ? handle.get_comms().get_rank() : 0; + int comm_size = distributed ? handle.get_comms().get_size() : 1; // Build exclusive global row offsets so sampled global row ids can be mapped to owning ranks. rmm::device_uvector rank_row_offsets(comm_size + 1, stream); diff --git a/cpp/src/randomforest/randomforest.cuh b/cpp/src/randomforest/randomforest.cuh index 8ce5992eef..b44c1d3c80 100644 --- a/cpp/src/randomforest/randomforest.cuh +++ b/cpp/src/randomforest/randomforest.cuh @@ -10,6 +10,7 @@ #include #include +#include #include #include #include @@ -90,7 +91,9 @@ class RowSampler { sample_weight_cdf_.begin()); } - if (sample_weight_ != nullptr) { + // Empty distributed partitions still pass a non-null pointer so every rank selects the same + // weighted objective type, but they have no local weight sum to validate. + if (sample_weight_ != nullptr && n_rows_ > 0) { sample_weight_sum_ = compute_sample_weight_sum(handle); ASSERT(sample_weight_sum_ > 0.0, "sample_weight values must contain at least one positive value"); @@ -114,6 +117,7 @@ class RowSampler { raft::common::nvtx::range fun_scope("bootstrapping row IDs @randomforest.cuh"); auto& selected_rows = selected_rows_[stream_id]; + if (n_rows_ == 0) { return selected_rows; } raft::resources stream_resources; raft::resource::set_cuda_stream(stream_resources, stream); @@ -234,14 +238,21 @@ class RandomForest { RF_params rf_params; // structure containing RF hyperparameters int rf_type; // 0 for classification 1 for regression - void error_checking(const T* input, L* predictions, int n_rows, int n_cols, bool predict) const + void error_checking(const T* input, + L* predictions, + int n_rows, + int n_cols, + bool predict, + bool allow_empty_local_rows = false) const { if (predict) { ASSERT(predictions != nullptr, "Error! User has not allocated memory for predictions."); } - ASSERT((n_rows > 0), "Invalid n_rows %d", n_rows); + ASSERT(allow_empty_local_rows ? (n_rows >= 0) : (n_rows > 0), "Invalid n_rows %d", n_rows); ASSERT((n_cols > 0), "Invalid n_cols %d", n_cols); + if (n_rows == 0) { return; } + bool input_is_dev_ptr = DT::is_dev_ptr(input); bool preds_is_dev_ptr = DT::is_dev_ptr(predictions); @@ -297,8 +308,10 @@ class RandomForest { bool input_row_major = false) { raft::common::nvtx::range fun_scope("RandomForest::fit @randomforest.cuh"); - this->error_checking(input, labels, n_rows, n_cols, false); - const raft::handle_t& handle = user_handle; + const raft::handle_t& handle = user_handle; + bool distributed = + raft::resource::comms_initialized(handle) && handle.get_comms().get_size() > 1; + this->error_checking(input, labels, n_rows, n_cols, false, distributed); std::int64_t const n_rows_i64 = n_rows; std::int64_t n_sampled_rows = 0; if (this->rf_params.bootstrap) { @@ -314,8 +327,11 @@ class RandomForest { n_sampled_rows = n_rows_i64; } int n_streams = this->rf_params.n_streams; + // Distributed tree builders issue collectives independently, so train them serially until + // the forest-level scheduler can impose a global collective order across concurrent trees. + if (distributed) { n_streams = 1; } ASSERT(static_cast(n_streams) <= handle.get_stream_pool_size(), - "rf_params.n_streams (=%d) should be <= raft::handle_t.n_streams (=%lu)", + "effective RF n_streams (=%d) should be <= raft::handle_t.n_streams (=%lu)", n_streams, handle.get_stream_pool_size()); diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 3bb0cc4325..307368fac4 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -229,6 +229,7 @@ if(BUILD_CUML_MG_TESTS) ConfigureTest( PREFIX MG NAME RF_QUANTILE_TEST mg/rf_quantile_test.cu MPI RAFT_DISTRIBUTED ML_INCLUDE ) + ConfigureTest(PREFIX MG NAME RF_TEST mg/rf_test.cu MPI RAFT_DISTRIBUTED ML_INCLUDE) else(MPI_CXX_FOUND) message("OpenMPI not found. Skipping MultiGPU tests '${CUML_MG_TEST_TARGET}'") endif() diff --git a/cpp/tests/mg/rf_test.cu b/cpp/tests/mg/rf_test.cu new file mode 100644 index 0000000000..69a464a599 --- /dev/null +++ b/cpp/tests/mg/rf_test.cu @@ -0,0 +1,681 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../prims/test_utils.h" +#include "test_opg_utils.h" + +#include + +#include +#include +#include + +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace ML { +namespace Test { +namespace opg { + +enum class PartitionKind { Contiguous, Strided, Imbalanced, EmptyNonRootRanks }; + +struct RfMgTestParams { + int n_rows; + int n_cols; + int n_trees; + float max_features; + int max_depth; + int max_leaves; + int max_n_bins; + int min_samples_leaf; + int min_samples_split; + float min_impurity_decrease; + int n_streams; + int handle_n_streams; + CRITERION split_criterion; + int seed; + int n_labels; + bool double_precision; + bool use_sample_weights; + PartitionKind partition_kind; +}; + +template +void expect_floating_values_equal(T distributed_value, T single_node_value, char const* field) +{ + static_assert(std::is_floating_point_v); + constexpr T absolute_tolerance = std::is_same_v ? T{1e-5} : T{1e-12}; + constexpr T relative_tolerance = std::is_same_v ? T{1e-5} : T{1e-10}; + auto scale = std::max(std::abs(distributed_value), std::abs(single_node_value)); + auto tolerance = absolute_tolerance + relative_tolerance * scale; + EXPECT_NEAR(distributed_value, single_node_value, tolerance) + << "Mismatched distributed RF " << field; +} + +template +void expect_forests_equal(RandomForestMetaData const& distributed_forest, + RandomForestMetaData const& single_node_forest) +{ + EXPECT_EQ(distributed_forest.n_features, single_node_forest.n_features); + + auto const& distributed_params = distributed_forest.rf_params; + auto const& single_node_params = single_node_forest.rf_params; + EXPECT_EQ(distributed_params.n_trees, single_node_params.n_trees); + EXPECT_EQ(distributed_params.bootstrap, single_node_params.bootstrap); + expect_floating_values_equal( + distributed_params.max_samples, single_node_params.max_samples, "max_samples"); + EXPECT_EQ(distributed_params.seed, single_node_params.seed); + + auto const& distributed_tree_params = distributed_params.tree_params; + auto const& single_node_tree_params = single_node_params.tree_params; + EXPECT_EQ(distributed_tree_params.max_depth, single_node_tree_params.max_depth); + EXPECT_EQ(distributed_tree_params.max_leaves, single_node_tree_params.max_leaves); + expect_floating_values_equal( + distributed_tree_params.max_features, single_node_tree_params.max_features, "max_features"); + EXPECT_EQ(distributed_tree_params.max_n_bins, single_node_tree_params.max_n_bins); + EXPECT_EQ(distributed_tree_params.min_samples_leaf, single_node_tree_params.min_samples_leaf); + EXPECT_EQ(distributed_tree_params.min_samples_split, single_node_tree_params.min_samples_split); + expect_floating_values_equal(distributed_tree_params.min_impurity_decrease, + single_node_tree_params.min_impurity_decrease, + "min_impurity_decrease"); + EXPECT_EQ(distributed_tree_params.split_criterion, single_node_tree_params.split_criterion); + EXPECT_EQ(distributed_tree_params.max_batch_size, single_node_tree_params.max_batch_size); + + ASSERT_EQ(distributed_forest.trees.size(), single_node_forest.trees.size()); + for (size_t tree_idx = 0; tree_idx < distributed_forest.trees.size(); ++tree_idx) { + SCOPED_TRACE(::testing::Message() << "tree_idx=" << tree_idx); + auto const& distributed_tree = distributed_forest.trees[tree_idx]; + auto const& single_node_tree = single_node_forest.trees[tree_idx]; + ASSERT_NE(distributed_tree, nullptr); + ASSERT_NE(single_node_tree, nullptr); + EXPECT_EQ(distributed_tree->treeid, single_node_tree->treeid); + EXPECT_EQ(distributed_tree->depth_counter, single_node_tree->depth_counter); + EXPECT_EQ(distributed_tree->leaf_counter, single_node_tree->leaf_counter); + EXPECT_EQ(distributed_tree->num_outputs, single_node_tree->num_outputs); + + ASSERT_EQ(distributed_tree->sparsetree.size(), single_node_tree->sparsetree.size()); + for (size_t node_idx = 0; node_idx < distributed_tree->sparsetree.size(); ++node_idx) { + SCOPED_TRACE(::testing::Message() << "node_idx=" << node_idx); + auto const& distributed_node = distributed_tree->sparsetree[node_idx]; + auto const& single_node_node = single_node_tree->sparsetree[node_idx]; + EXPECT_EQ(distributed_node.ColumnId(), single_node_node.ColumnId()); + expect_floating_values_equal( + distributed_node.QueryValue(), single_node_node.QueryValue(), "split threshold"); + expect_floating_values_equal( + distributed_node.BestMetric(), single_node_node.BestMetric(), "split metric"); + EXPECT_EQ(distributed_node.LeftChildId(), single_node_node.LeftChildId()); + EXPECT_EQ(distributed_node.InstanceCount(), single_node_node.InstanceCount()); + EXPECT_EQ(distributed_node.IsLeaf(), single_node_node.IsLeaf()); + } + + ASSERT_EQ(distributed_tree->vector_leaf.size(), single_node_tree->vector_leaf.size()); + for (size_t value_idx = 0; value_idx < distributed_tree->vector_leaf.size(); ++value_idx) { + SCOPED_TRACE(::testing::Message() << "leaf_value_idx=" << value_idx); + expect_floating_values_equal(distributed_tree->vector_leaf[value_idx], + single_node_tree->vector_leaf[value_idx], + "leaf value"); + } + } +} + +std::vector local_rows_for_rank(int n_rows, int rank, int size, PartitionKind kind) +{ + std::vector rows; + if (kind == PartitionKind::Strided) { + for (int row = rank; row < n_rows; row += size) { + rows.push_back(row); + } + return rows; + } + + std::vector counts(size, n_rows / size); + for (int i = 0; i < n_rows % size; ++i) { + counts[i]++; + } + if (kind == PartitionKind::Imbalanced && size > 1) { + counts.assign(size, 0); + counts[0] = std::max(1, (n_rows * 3) / 4); + int remaining = n_rows - counts[0]; + for (int i = 1; i < size; ++i) { + counts[i] = remaining / (size - 1); + } + for (int i = 1; i <= remaining % (size - 1); ++i) { + counts[i]++; + } + } else if (kind == PartitionKind::EmptyNonRootRanks && size > 1) { + counts.assign(size, 0); + counts[0] = n_rows; + } + + int begin = std::accumulate(counts.begin(), counts.begin() + rank, 0); + rows.resize(counts[rank]); + std::iota(rows.begin(), rows.end(), begin); + rows.erase(std::remove_if(rows.begin(), rows.end(), [=](int row) { return row >= n_rows; }), + rows.end()); + return rows; +} + +template +void make_local_dataset(RfMgTestParams const& params, + std::vector const& rows, + std::vector& X, + std::vector& y, + std::vector& sample_weights) +{ + X.resize(rows.size() * params.n_cols); + y.resize(rows.size()); + sample_weights.resize(params.use_sample_weights ? rows.size() : 0); + for (size_t i = 0; i < rows.size(); ++i) { + int global_row = rows[i]; + DataT signal = static_cast((global_row % 97) - 48); + for (int col = 0; col < params.n_cols; ++col) { + DataT feature = signal * static_cast(col + 1); + feature += static_cast(((global_row + 13 * col + params.seed) % 11) - 5) / + static_cast(10); + X[static_cast(col) * rows.size() + i] = feature; + } + if constexpr (std::is_integral_v) { + y[i] = (signal >= DataT(0)) ? 1 : 0; + if (params.n_labels > 2 && global_row % 17 == 0) { y[i] = 2; } + } else { + y[i] = signal * DataT(0.5) + static_cast((global_row % 7) - 3); + } + if (params.use_sample_weights) { sample_weights[i] = global_row % 2 == 0 ? 0.8 : 0.6; } + } +} + +std::vector global_rows_in_rank_order(RfMgTestParams const& params, int size) +{ + std::vector rows; + rows.reserve(params.n_rows); + for (int rank = 0; rank < size; ++rank) { + auto rank_rows = local_rows_for_rank(params.n_rows, rank, size, params.partition_kind); + rows.insert(rows.end(), rank_rows.begin(), rank_rows.end()); + } + return rows; +} + +template +void expect_global_tree_counts(RandomForestMetaData const& forest, int n_rows) +{ + for (auto const& tree : forest.trees) { + ASSERT_FALSE(tree->sparsetree.empty()); + EXPECT_EQ(tree->sparsetree.front().InstanceCount(), n_rows); + for (auto const& node : tree->sparsetree) { + if (!node.IsLeaf()) { + ASSERT_GE(node.LeftChildId(), 0); + ASSERT_GE(node.RightChildId(), 0); + ASSERT_LT(static_cast(node.LeftChildId()), tree->sparsetree.size()); + ASSERT_LT(static_cast(node.RightChildId()), tree->sparsetree.size()); + auto left_count = tree->sparsetree[node.LeftChildId()].InstanceCount(); + auto right_count = tree->sparsetree[node.RightChildId()].InstanceCount(); + EXPECT_EQ(left_count + right_count, node.InstanceCount()); + } + } + } +} + +template +void expect_tree_limits(RandomForestMetaData const& forest, RfMgTestParams const& params) +{ + EXPECT_EQ(forest.trees.size(), params.n_trees); + for (auto const& tree : forest.trees) { + EXPECT_LE(tree->depth_counter, params.max_depth); + if (params.max_leaves > 0) { EXPECT_LE(tree->leaf_counter, params.max_leaves); } + for (auto const& node : tree->sparsetree) { + if (!node.IsLeaf()) { EXPECT_GT(node.BestMetric(), params.min_impurity_decrease); } + } + } +} + +void initialize_mpi_once() +{ + int mpi_initialized = 0; + MPI_Initialized(&mpi_initialized); + if (!mpi_initialized) { MPI_Init(nullptr, nullptr); } +} + +void get_mpi_local_rank_size(int& local_rank, int& local_size) +{ + MPI_Comm local_comm{}; + MPI_Comm_split_type(MPI_COMM_WORLD, MPI_COMM_TYPE_SHARED, 0, MPI_INFO_NULL, &local_comm); + MPI_Comm_rank(local_comm, &local_rank); + MPI_Comm_size(local_comm, &local_size); + MPI_Comm_free(&local_comm); +} + +template +class RfMgPropertyTestImpl { + public: + explicit RfMgPropertyTestImpl(RfMgTestParams const& params) : params(params) + { + initialize_mpi_once(); + int rank = 0; + int size = 1; + MPI_Comm_rank(MPI_COMM_WORLD, &rank); + MPI_Comm_size(MPI_COMM_WORLD, &size); + + int local_rank = 0; + int local_size = 1; + get_mpi_local_rank_size(local_rank, local_size); + + int n_gpus = 0; + RAFT_CUDA_TRY(cudaGetDeviceCount(&n_gpus)); + int insufficient_local_gpus = n_gpus < local_size; + int any_insufficient_gpus = 0; + MPI_Allreduce( + &insufficient_local_gpus, &any_insufficient_gpus, 1, MPI_INT, MPI_MAX, MPI_COMM_WORLD); + if (any_insufficient_gpus) { + if (insufficient_local_gpus) { + ADD_FAILURE() << "Number of GPUs is smaller than local MPI ranks: ngpus=" << n_gpus + << ", local_ranks=" << local_size; + } + return; + } + RAFT_CUDA_TRY(cudaSetDevice(local_rank)); + + auto stream_pool = std::make_shared(params.handle_n_streams); + raft::handle_t handle(rmm::cuda_stream_per_thread, stream_pool); + raft::comms::initialize_mpi_comms(&handle, MPI_COMM_WORLD); + + auto local_rows = local_rows_for_rank(params.n_rows, rank, size, params.partition_kind); + std::vector h_X; + std::vector h_y; + std::vector h_sample_weights; + make_local_dataset(params, local_rows, h_X, h_y, h_sample_weights); + + rmm::device_uvector X(h_X.size(), handle.get_stream()); + rmm::device_uvector y(h_y.size(), handle.get_stream()); + // A non-null pointer keeps empty ranks on the weighted objective path. + auto sample_weight_buffer_size = + params.use_sample_weights ? std::max(size_t{1}, h_sample_weights.size()) : size_t{0}; + rmm::device_uvector sample_weights(sample_weight_buffer_size, handle.get_stream()); + raft::update_device(X.data(), h_X.data(), h_X.size(), handle.get_stream()); + raft::update_device(y.data(), h_y.data(), h_y.size(), handle.get_stream()); + raft::update_device( + sample_weights.data(), h_sample_weights.data(), h_sample_weights.size(), handle.get_stream()); + auto sample_weight_ptr = params.use_sample_weights ? sample_weights.data() : nullptr; + + auto rf_params = set_rf_params(params.max_depth, + params.max_leaves, + params.max_features, + params.max_n_bins, + params.min_samples_leaf, + params.min_samples_split, + params.min_impurity_decrease, + false, + params.n_trees, + 1.0f, + params.seed, + params.split_criterion, + params.n_streams, + 128); + + RandomForestMetaData distributed_forest; + if constexpr (std::is_integral_v) { + fit(handle, + &distributed_forest, + X.data(), + static_cast(local_rows.size()), + params.n_cols, + y.data(), + params.n_labels, + rf_params, + rapids_logger::level_enum::info, + nullptr, + sample_weight_ptr); + } else { + fit(handle, + &distributed_forest, + X.data(), + static_cast(local_rows.size()), + params.n_cols, + y.data(), + rf_params, + rapids_logger::level_enum::info, + nullptr, + sample_weight_ptr); + } + + expect_global_tree_counts(distributed_forest, params.n_rows); + expect_tree_limits(distributed_forest, params); + + // The distributed quantile sampler assigns global row ids in rank-major order. Reconstruct the + // single-node input in the same order so both algorithms receive the identical training data. + auto global_rows = global_rows_in_rank_order(params, size); + if (global_rows.size() != static_cast(params.n_rows)) { + ADD_FAILURE() << "Reconstructed " << global_rows.size() << " global rows, expected " + << params.n_rows; + return; + } + std::vector h_global_X; + std::vector h_global_y; + std::vector h_global_sample_weights; + make_local_dataset( + params, global_rows, h_global_X, h_global_y, h_global_sample_weights); + + auto single_node_stream_pool = std::make_shared(params.handle_n_streams); + raft::handle_t single_node_handle(rmm::cuda_stream_per_thread, single_node_stream_pool); + rmm::device_uvector global_X(h_global_X.size(), single_node_handle.get_stream()); + rmm::device_uvector global_y(h_global_y.size(), single_node_handle.get_stream()); + rmm::device_uvector global_sample_weights(h_global_sample_weights.size(), + single_node_handle.get_stream()); + raft::update_device( + global_X.data(), h_global_X.data(), h_global_X.size(), single_node_handle.get_stream()); + raft::update_device( + global_y.data(), h_global_y.data(), h_global_y.size(), single_node_handle.get_stream()); + raft::update_device(global_sample_weights.data(), + h_global_sample_weights.data(), + h_global_sample_weights.size(), + single_node_handle.get_stream()); + auto global_sample_weight_ptr = + params.use_sample_weights ? global_sample_weights.data() : nullptr; + + auto single_node_rf_params = rf_params; + single_node_rf_params.n_streams = 1; + RandomForestMetaData single_node_forest; + if constexpr (std::is_integral_v) { + fit(single_node_handle, + &single_node_forest, + global_X.data(), + params.n_rows, + params.n_cols, + global_y.data(), + params.n_labels, + single_node_rf_params, + rapids_logger::level_enum::info, + nullptr, + global_sample_weight_ptr); + } else { + fit(single_node_handle, + &single_node_forest, + global_X.data(), + params.n_rows, + params.n_cols, + global_y.data(), + single_node_rf_params, + rapids_logger::level_enum::info, + nullptr, + global_sample_weight_ptr); + } + + expect_forests_equal(distributed_forest, single_node_forest); + } + + private: + RfMgTestParams params; +}; + +class RfMgPropertyTest : public ::testing::TestWithParam { + public: + void SetUp() override + { + auto params = GetParam(); + bool is_regression = params.split_criterion != GINI && params.split_criterion != ENTROPY; + if (params.double_precision) { + if (is_regression) { + RfMgPropertyTestImpl test(params); + } else { + RfMgPropertyTestImpl test(params); + } + } else { + if (is_regression) { + RfMgPropertyTestImpl test(params); + } else { + RfMgPropertyTestImpl test(params); + } + } + } +}; + +TEST_P(RfMgPropertyTest, DistributedProperties) {} + +constexpr auto UNLIMITED_DEPTH = std::numeric_limits::max(); + +std::vector inputs = { + {128, + 4, + 1, + 1.0f, + 3, + -1, + 16, + 1, + 2, + 0.0f, + 1, + 1, + GINI, + 7, + 2, + false, + false, + PartitionKind::Contiguous}, + {128, + 4, + 3, + 0.5f, + 4, + 16, + 32, + 1, + 2, + 0.0f, + 4, + 4, + ENTROPY, + 11, + 2, + false, + false, + PartitionKind::Strided}, + {192, + 6, + 1, + 1.0f, + 5, + -1, + 32, + 2, + 4, + 0.0f, + 1, + 1, + MSE, + 13, + 2, + false, + false, + PartitionKind::Imbalanced}, + {96, 3, 2, 1.0f, 4, 8, 8, 1, 2, 0.0f, 1, 1, GINI, 17, 2, true, false, PartitionKind::Imbalanced}, + {144, 5, 2, 0.8f, 4, -1, 16, 1, 2, 0.0f, 1, 1, GINI, 31, 3, false, false, PartitionKind::Strided}, + {160, + 5, + 1, + 0.8f, + 4, + -1, + 16, + 1, + 2, + 0.0f, + 1, + 1, + MSE, + 19, + 2, + true, + false, + PartitionKind::Contiguous}, + {256, + 4, + 5, + 1.0f, + 5, + -1, + 16, + 1, + 2, + 0.0f, + 1, + 1, + POISSON, + 7, + 1, + true, + false, + PartitionKind::Strided}, + {256, + 4, + 2, + 0.8f, + 5, + -1, + 16, + 1, + 2, + 0.0f, + 1, + 1, + GAMMA, + 7, + 1, + true, + false, + PartitionKind::Contiguous}, + {256, + 4, + 2, + 0.8f, + 5, + -1, + 16, + 1, + 2, + 0.0f, + 1, + 1, + INVERSE_GAUSSIAN, + 7, + 1, + true, + false, + PartitionKind::Contiguous}, + {256, + 4, + 5, + 1.0f, + UNLIMITED_DEPTH, + -1, + 16, + 1, + 2, + 0.0f, + 1, + 1, + GINI, + 7, + 2, + false, + false, + PartitionKind::Imbalanced}, + {160, + 4, + 5, + 1.0f, + UNLIMITED_DEPTH, + -1, + 16, + 1, + 2, + 0.0f, + 1, + 1, + ENTROPY, + 7, + 2, + false, + false, + PartitionKind::Imbalanced}, + {64, + 4, + 2, + 1.0f, + 4, + -1, + 16, + 1, + 2, + 0.0f, + 3, + 3, + GINI, + 23, + 2, + false, + false, + PartitionKind::EmptyNonRootRanks}, + {80, + 5, + 2, + 0.8f, + 4, + -1, + 16, + 1, + 2, + 0.0f, + 3, + 3, + MSE, + 29, + 2, + false, + false, + PartitionKind::EmptyNonRootRanks}, + {80, + 5, + 2, + 0.8f, + 4, + -1, + 16, + 1, + 2, + 0.0f, + 3, + 3, + MSE, + 29, + 2, + false, + true, + PartitionKind::EmptyNonRootRanks}}; + +INSTANTIATE_TEST_SUITE_P(RfTests, RfMgPropertyTest, ::testing::ValuesIn(inputs)); + +} // namespace opg +} // namespace Test +} // namespace ML + +int main(int argc, char** argv) +{ + ::testing::InitGoogleTest(&argc, argv); + ::testing::AddGlobalTestEnvironment(new MLCommon::Test::opg::MPIEnvironment()); + + return RUN_ALL_TESTS(); +} diff --git a/cpp/tests/sg/rf_test.cu b/cpp/tests/sg/rf_test.cu index 12ac6f8ccb..17cea6caa2 100644 --- a/cpp/tests/sg/rf_test.cu +++ b/cpp/tests/sg/rf_test.cu @@ -951,6 +951,21 @@ TEST(RfTests, IntegerOverflow) handle.sync_stream_pool(); } +TEST(RfTests, EmptyGlobalRowsRejected) +{ + thrust::device_vector X(1); + thrust::device_vector y(1); + auto forest = std::make_shared>(); + auto forest_ptr = forest.get(); + auto stream_pool = std::make_shared(1); + raft::handle_t handle(rmm::cuda_stream_per_thread, stream_pool); + RF_params rf_params = + set_rf_params(3, 100, 1.0, 16, 1, 2, 0.0, false, 1, 1.0, 0, CRITERION::MSE, 1, 128); + + EXPECT_THROW(fit(handle, forest_ptr, X.data().get(), 0, 1, y.data().get(), rf_params), + raft::exception); +} + TEST(RfTests, HighClassCountSplitHistogramFallsBackToGlobalMemory) { constexpr std::size_t n_rows = 640; From 74119ad8c87daa1c36fd016a15e7c1a7d7fafbac Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Tue, 11 Aug 2026 05:31:17 +0000 Subject: [PATCH 2/2] Update Dask RF to use the new distributed algo --- python/cuml/cuml/dask/ensemble/base.py | 297 ++++-------------- .../dask/ensemble/randomforestclassifier.py | 100 +----- .../dask/ensemble/randomforestregressor.py | 81 +---- .../cuml/ensemble/randomforest_common.pyx | 17 +- .../cuml/ensemble/randomforestclassifier.py | 3 +- .../tests/dask/test_dask_random_forest.py | 199 ++++-------- 6 files changed, 160 insertions(+), 537 deletions(-) diff --git a/python/cuml/cuml/dask/ensemble/base.py b/python/cuml/cuml/dask/ensemble/base.py index 65b01eb67a..f21df3b11b 100644 --- a/python/cuml/cuml/dask/ensemble/base.py +++ b/python/cuml/cuml/dask/ensemble/base.py @@ -2,18 +2,11 @@ # SPDX-License-Identifier: Apache-2.0 # -import math -import warnings -from collections.abc import Iterable +from dask.distributed import get_worker +from raft_dask.common.comms import Comms, get_raft_comm_state -import cupy as cp -import dask -import numpy as np -import treelite -from dask.distributed import Future - -from cuml import using_output_type from cuml.dask._compat import DASK_2025_4_0 +from cuml.dask.common.base import mnmg_import from cuml.dask.common.input_utils import DistributedDataHandler, concatenate from cuml.dask.common.utils import get_client, wait_and_raise_from_futures @@ -34,7 +27,6 @@ def _create_model( workers, n_estimators, base_seed, - ignore_empty_partitions, **kwargs, ): self.client = get_client(client) @@ -46,180 +38,76 @@ def _create_model( ) self.workers = workers self._set_internal_model(None) - self.active_workers = list() - self.ignore_empty_partitions = ignore_empty_partitions self.n_estimators = n_estimators - - self.n_estimators_per_worker = self._estimators_per_worker( - n_estimators - ) - if base_seed is None: - base_seed = 0 - seeds = [base_seed] - for i in range(1, len(self.n_estimators_per_worker)): - sd = self.n_estimators_per_worker[i - 1] + seeds[i - 1] - seeds.append(sd) + self.n_streams = kwargs.get("n_streams", 4) self.rfs = { worker: self.client.submit( model_func, - n_estimators=self.n_estimators_per_worker[n], - random_state=seeds[n], + n_estimators=self.n_estimators, + random_state=base_seed, **kwargs, pure=False, workers=[worker], ) - for n, worker in enumerate(self.workers) + for worker in self.workers } wait_and_raise_from_futures(list(self.rfs.values())) - def _estimators_per_worker(self, n_estimators): - n_workers = len(self.workers) - if n_estimators < n_workers: - raise ValueError( - "n_estimators cannot be lower than number of dask workers." - ) - - n_est_per_worker = math.floor(n_estimators / n_workers) - n_estimators_per_worker = [n_est_per_worker for i in range(n_workers)] - remaining_est = n_estimators - (n_est_per_worker * n_workers) - for i in range(remaining_est): - n_estimators_per_worker[i] = n_estimators_per_worker[i] + 1 - return n_estimators_per_worker - - def _fit(self, model, dataset, broadcast_data): + def _fit(self, model, dataset, classes=None): data = DistributedDataHandler.create(dataset, client=self.client) - self.active_workers = data.workers self.datatype = data.datatype - labels = self.client.persist(dataset[1]) - if self.datatype == "cudf": - self.num_classes = len(labels.unique()) - else: - self.num_classes = len(dask.array.unique(labels).compute()) + unknown_workers = set(data.workers).difference(model) + if unknown_workers: + raise ValueError( + "Training data was placed on workers that were not selected " + f"for this estimator: {sorted(unknown_workers)}" + ) - combined_data = ( - list(map(lambda x: x[1], data.gpu_futures)) - if broadcast_data - else None + total_rows = sum(total for _, total in data._worker_sizes.values()) + comms = Comms( + comms_p2p=False, + client=self.client, + streams_per_handle=self.n_streams, ) + comms.init(workers=data.workers) - futures = list() - for idx, (worker, worker_data) in enumerate( - data.worker_to_parts.items() - ): - futures.append( - self.client.submit( + futures = [] + try: + for worker, worker_data in data.worker_to_parts.items(): + future = self.client.submit( _func_fit, + comms.sessionId, model[worker], - combined_data if broadcast_data else worker_data, + worker_data, + total_rows, + classes, workers=[worker], pure=False, ) - ) + futures.append(future) + self.rfs[worker] = future - self.n_active_estimators_per_worker = [] - for worker in data.worker_to_parts.keys(): - n = self.workers.index(worker) - n_est = self.n_estimators_per_worker[n] - self.n_active_estimators_per_worker.append(n_est) + wait_and_raise_from_futures(futures) + finally: + comms.destroy() - if len(self.workers) > len(self.active_workers): - if self.ignore_empty_partitions: - curent_estimators = ( - self.n_estimators - / len(self.workers) - * len(self.active_workers) - ) - warn_text = ( - f"Data was not split among all workers " - f"using only {self.active_workers} workers to fit." - f"This will only train {curent_estimators}" - f" estimators instead of the requested " - f"{self.n_estimators}" - ) - warnings.warn(warn_text) - else: - raise ValueError( - "Data was not split among all workers. " - "Re-run the code or " - "use ignore_empty_partitions=True" - " while creating model" - ) - wait_and_raise_from_futures(futures) + # Every distributed rank owns the same complete forest. Keep one + # worker future as the canonical model for inference and serialization. + self._set_internal_model(futures[0]) return self - def _concat_treelite_models(self): - """ - Convert the cuML Random Forest model present in different workers to - the treelite format and then concatenate the different treelite models - to create a single model. The concatenated model is then converted to - bytes format. - """ - model_serialized_futures = list() - for w in self.active_workers: - model_serialized_futures.append( - dask.delayed(_serialize_treelite_bytes)(self.rfs[w]) - ) - mod_bytes = self.client.compute(model_serialized_futures, sync=True) - last_worker = w - model = self.rfs[last_worker].result() - tl_model_objs = [ - treelite.Model.deserialize_bytes(indiv_worker_model_bytes) - for indiv_worker_model_bytes in mod_bytes - ] - concatenated_model = treelite.Model.concatenate(tl_model_objs) - model._treelite_model_bytes = concatenated_model.serialize_bytes() - model._fil_model = None - return model - - def _partial_inference(self, X, op_type, delayed, **kwargs): + def _predict_using_nvforest(self, X, delayed, **kwargs): data = DistributedDataHandler.create(X, client=self.client) - combined_data = list(map(lambda x: x[1], data.gpu_futures)) - - if op_type == "classification": - func = _func_predict_proba_partial - shape = (X.shape[0], 1, self.num_classes) - else: - shape = (X.shape[0], 1) - func = _func_predict_partial - - meta = cp.zeros((0,) * len(shape), dtype=cp.float32) - - partial_infs = list() - for worker in self.active_workers: - partial_infs.append( - self.client.submit( - func, - self.rfs[worker], - combined_data, - **kwargs, - workers=[worker], - pure=False, - ) - ) - - objs = [ - dask.array.from_delayed(partial_inf, shape=shape, meta=meta) - for partial_inf in partial_infs - ] - result = dask.array.concatenate(objs, axis=1) - return result - - def _predict_using_fil(self, X, delayed, **kwargs): - if self._get_internal_model() is None: - self._set_internal_model(self._concat_treelite_models()) - data = DistributedDataHandler.create(X, client=self.client) - if self._get_internal_model() is None: - self._set_internal_model(self._concat_treelite_models()) return self._predict( X, delayed=delayed, output_collection_type=data.datatype, **kwargs ) def _get_params(self, deep): model_params = list() - for idx, worker in enumerate(self.workers): + for worker in self.workers: model_params.append( self.client.submit( _func_get_params, self.rfs[worker], deep, workers=[worker] @@ -229,8 +117,10 @@ def _get_params(self, deep): return params_of_each_model def _set_params(self, **params): + if "n_streams" in params: + self.n_streams = params["n_streams"] model_params = list() - for idx, worker in enumerate(self.workers): + for worker in self.workers: model_params.append( self.client.submit( _func_set_params, @@ -242,96 +132,23 @@ def _set_params(self, **params): wait_and_raise_from_futures(model_params) return self - def get_combined_model(self): - """ - Return single-GPU model for serialization. - - Returns - ------- - - model : Trained single-GPU model or None if the model has not - yet been trained. - """ - - # set internal model if it hasn't been accessed before - if self._get_internal_model() is None: - self._set_internal_model(self._concat_treelite_models()) - - internal_model = self._check_internal_model(self._get_internal_model()) - - if isinstance(self.internal_model, Iterable): - # This function needs to return a single instance of cuml.Base, - # even if the class is just a composite. - raise ValueError( - "Expected a single instance of cuml.Base " - "but got %s instead." % type(self.internal_model) - ) - - elif isinstance(self.internal_model, Future): - internal_model = self.internal_model.result() - - return internal_model - - def _get_workers_weights(self) -> cp.ndarray: - workers_weights = np.array(self.n_active_estimators_per_worker) - workers_weights = workers_weights[workers_weights != 0] - workers_weights = workers_weights / workers_weights.sum() - workers_weights = cp.array(workers_weights) - return workers_weights - - def apply_reduction(self, reduce, partial_infs, datatype, delayed): - """ - Reduces the partial inferences to obtain the final result. The workers - didn't have the same number of trees to form their predictions. To - correct for this worker's predictions are weighted differently during - reduction. - """ - workers_weights = self._get_workers_weights() - unique_classes = ( - None - if not hasattr(self, "unique_classes") - else self.unique_classes - ) - delayed_local_array = dask.delayed(reduce)( - partial_infs, workers_weights, unique_classes - ) - delayed_res = dask.array.from_delayed( - delayed_local_array, shape=(np.nan, np.nan), dtype=np.float32 - ) - if delayed: - return delayed_res - else: - return delayed_res.persist() - -def _func_fit(model, input_data): +@mnmg_import +def _func_fit(session_id, model, input_data, total_rows, classes): + handle = get_raft_comm_state(session_id, get_worker())["handle"] X = concatenate([item[0] for item in input_data]) y = concatenate([item[1] for item in input_data]) - return model.fit(X, y) - - -def _func_predict_partial(model, input_data, **kwargs): - """ - Whole dataset inference with part of the model (trees at disposal locally). - Transfer dataset instead of model. Interesting when model is larger - than dataset. - """ - X = concatenate(input_data) - with using_output_type("cupy"): - prediction = model.predict(X, **kwargs) - return cp.expand_dims(prediction, axis=1) - - -def _func_predict_proba_partial(model, input_data, **kwargs): - """ - Whole dataset inference with part of the model (trees at disposal locally). - Transfer dataset instead of model. Interesting when model is larger - than dataset. - """ - X = concatenate(input_data) - with using_output_type("cupy"): - prediction = model.predict_proba(X, **kwargs) - return cp.expand_dims(prediction, axis=1) + model._raft_handle = handle + model._distributed_n_rows = total_rows + if classes is not None: + model._distributed_classes = classes + try: + return model.fit(X, y) + finally: + del model._raft_handle + del model._distributed_n_rows + if classes is not None: + del model._distributed_classes def _func_get_params(model, deep): @@ -340,7 +157,3 @@ def _func_get_params(model, deep): def _func_set_params(model, **params): return model.set_params(**params) - - -def _serialize_treelite_bytes(model): - return model._treelite_model_bytes diff --git a/python/cuml/cuml/dask/ensemble/randomforestclassifier.py b/python/cuml/cuml/dask/ensemble/randomforestclassifier.py index 97b9aca4d7..07425a7d13 100755 --- a/python/cuml/cuml/dask/ensemble/randomforestclassifier.py +++ b/python/cuml/cuml/dask/ensemble/randomforestclassifier.py @@ -26,30 +26,13 @@ class RandomForestClassifier( classifiers in an ensemble. This uses Dask to partition data over multiple GPUs (possibly on different nodes). - This implementation makes the following assumptions: - * The set of Dask workers used between instantiation, fit, \ - and predict are all consistent - * Training data comes in the form of cuDF dataframes or Dask Arrays \ - distributed so that each worker has at least one partition. - - The distributed algorithm uses an *embarrassingly-parallel* - approach. For a forest with `N` trees being built on `w` workers, each - worker simply builds `N/w` trees on the data it has available - locally. In many cases, partitioning the data so that each worker - builds trees on a subset of the total dataset works well, but - it generally requires the data to be well-shuffled in advance. - Alternatively, callers can replicate all of the data across - workers so that ``rf.fit`` receives `w` partitions, each containing the - same data. This would produce results approximately identical to - single-GPU fitting. - - Please check the single-GPU implementation of Random Forest - classifier for more information about the underlying algorithm. + During fitting, all workers that hold training rows collectively build the + same forest from the complete distributed dataset. Parameters ---------- n_estimators : int (default = 100) - total number of trees in the forest (not per-worker) + total number of trees in the forest split_criterion : int or string (default = ``0`` (``'gini'``)) The criterion used to split nodes.\n * ``0`` or ``'gini'`` for gini impurity @@ -107,20 +90,15 @@ class RandomForestClassifier( and ``ceil(min_samples_split * n_rows)`` is the minimum number of samples for each split. - n_streams : int (default = 4 ) - Number of parallel streams used for forest building + n_streams : int (default = 4) + Number of parallel streams requested for forest building. Distributed + training currently builds trees serially to preserve collective order. workers : optional, list of strings Dask addresses of workers to use for computation. If None, all available Dask workers will be used. random_state : int (default = None) Seed for the random number generator. Unseeded by default. - ignore_empty_partitions: Boolean (default = False) - Specify behavior when a worker does not hold any data - while splitting. When True, it returns the results from workers - with data (the number of trained estimators will be less than - n_estimators) When False, throws a RuntimeError. - Examples -------- For usage examples, please see the RAPIDS notebooks repository: @@ -135,7 +113,6 @@ def __init__( verbose=False, n_estimators=100, random_state=None, - ignore_empty_partitions=False, **kwargs, ): super().__init__(client=client, verbose=verbose, **kwargs) @@ -145,7 +122,6 @@ def __init__( workers=workers, n_estimators=n_estimators, base_seed=random_state, - ignore_empty_partitions=ignore_empty_partitions, **kwargs, ) @@ -155,12 +131,11 @@ def _construct_rf(n_estimators, random_state, **kwargs): n_estimators=n_estimators, random_state=random_state, **kwargs ) - def fit(self, X, y, broadcast_data=False): + def fit(self, X, y): """ Fit the input data with a Random Forest classifier - IMPORTANT: X is expected to be partitioned with at least one partition - on each Dask worker being used by the forest (self.workers). + Only workers holding one or more training rows participate in fitting. If a worker has multiple data partitions, they will be concatenated before fitting, which will lead to additional memory usage. To minimize @@ -196,27 +171,18 @@ def fit(self, X, y, broadcast_data=False): y : Dask cuDF dataframe or CuPy backed Dask Array (n_rows, 1) Labels of training examples. **y must be partitioned the same way as X** - broadcast_data : bool, optional (default = False) - When set to True, the whole dataset is broadcasted - to train the workers, otherwise each worker - is trained on its partition """ - # Handle both Dask Arrays and Dask Series/DataFrames if isinstance(y, dask.array.Array): - # For Dask Arrays, use dask.array.unique unique_vals = dask.array.unique(y).compute() - self.unique_classes = cp.sort(cp.asarray(unique_vals)) else: - # For Dask Series/DataFrames, use .unique() method - self.unique_classes = cp.asarray( - y.unique().compute().sort_values(ignore_index=True) - ) - self.num_classes = len(self.unique_classes) + unique_vals = y.unique().compute().sort_values(ignore_index=True) + classes = cp.asnumpy(cp.sort(cp.asarray(unique_vals))) + self.classes_ = classes self._set_internal_model(None) self._fit( model=self.rfs, dataset=(X, y), - broadcast_data=broadcast_data, + classes=classes, ) return self @@ -228,7 +194,6 @@ def predict( default_chunk_size=None, align_bytes=None, delayed=True, - broadcast_data=False, ): """ Predicts the labels for X. @@ -241,7 +206,7 @@ def predict( threshold : float (default = 0.5) Threshold used for classification. layout : string (default = 'depth_first') - Specifies the in-memory layout of nodes in FIL forests. Options: + Specifies the in-memory layout of nodes in nvForest models. Options: 'depth_first', 'layered', 'breadth_first'. default_chunk_size : int, optional (default = None) Determines how batches are further subdivided for parallel processing. @@ -255,28 +220,13 @@ def predict( delayed : bool (default = True) Whether to do a lazy prediction (and return Delayed objects) or an eagerly executed one. - broadcast_data : bool (default = False) - If False, the trees are merged in a single model before the workers - perform inference on their share of the prediction workload. - When True, trees aren't merged. Instead each worker infers on the - whole prediction workload using its available trees. The results are - reduced on the client. May be advantageous when the model is larger - than the data used for inference. Returns ------- y : Dask cuDF dataframe or CuPy backed Dask Array (n_rows, 1) The predicted class labels. """ - if broadcast_data: - return self.partial_inference( - X, - layout=layout, - default_chunk_size=default_chunk_size, - align_bytes=align_bytes, - delayed=delayed, - ) - return self._predict_using_fil( + return self._predict_using_nvforest( X, threshold=threshold, layout=layout, @@ -285,26 +235,6 @@ def predict( delayed=delayed, ) - def partial_inference(self, X, delayed, **kwargs): - partial_infs = self._partial_inference( - X=X, op_type="classification", delayed=delayed, **kwargs - ) - worker_weights = self._get_workers_weights() - merged_votes = dask.array.average( - partial_infs, axis=1, weights=worker_weights - ) - pred_class_indices = merged_votes.argmax(axis=1) - unique_classes = self.unique_classes - - pred_class = pred_class_indices.map_blocks( - lambda x: unique_classes[x], - meta=unique_classes[:0], - ) - if delayed: - return pred_class - else: - return pred_class.persist() - def predict_proba(self, X, delayed=True, **kwargs): """ Predicts the probability of each class for X. @@ -326,8 +256,6 @@ def predict_proba(self, X, delayed=True, **kwargs): ------- y : Dask cuDF dataframe or CuPy backed Dask Array (n_rows, n_classes) """ - if self._get_internal_model() is None: - self._set_internal_model(self._concat_treelite_models()) data = DistributedDataHandler.create(X, client=self.client) return self._predict_proba( X, delayed, output_collection_type=data.datatype, **kwargs diff --git a/python/cuml/cuml/dask/ensemble/randomforestregressor.py b/python/cuml/cuml/dask/ensemble/randomforestregressor.py index c4ccdea294..7eb7f627e4 100755 --- a/python/cuml/cuml/dask/ensemble/randomforestregressor.py +++ b/python/cuml/cuml/dask/ensemble/randomforestregressor.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # -import dask.array from cuml.dask.common.base import BaseEstimator, DelayedPredictionMixin from cuml.dask.ensemble.base import BaseRandomForestModel @@ -17,30 +16,13 @@ class RandomForestRegressor( regressors in an ensemble. This uses Dask to partition data over multiple GPUs (possibly on different nodes). - This implementation makes the following assumptions: - * The set of Dask workers used between instantiation, fit, - and predict are all consistent - * Training data comes in the form of cuDF dataframes or Dask Arrays - distributed so that each worker has at least one partition. - - The distributed algorithm uses an *embarrassingly-parallel* - approach. For a forest with `N` trees being built on `w` workers, each - worker simply builds `N/w` trees on the data it has available - locally. In many cases, partitioning the data so that each worker - builds trees on a subset of the total dataset works well, but - it generally requires the data to be well-shuffled in advance. - Alternatively, callers can replicate all of the data across - workers so that ``rf.fit`` receives `w` partitions, each containing the - same data. This would produce results approximately identical to - single-GPU fitting. - - Please check the single-GPU implementation of Random Forest - regressor for more information about the underlying algorithm. + During fitting, all workers that hold training rows collectively build the + same forest from the complete distributed dataset. Parameters ---------- n_estimators : int (default = 100) - total number of trees in the forest (not per-worker) + total number of trees in the forest split_criterion : int or string (default = ``2`` (``'mse'``)) The criterion used to split nodes.\n * ``0`` or ``'gini'`` for gini impurity @@ -93,20 +75,15 @@ class RandomForestRegressor( * If type ``float``, then ``min_samples_split`` represents a fraction and ``ceil(min_samples_split * n_rows)`` is the minimum number of samples for each split. - n_streams : int (default = 4 ) - Number of parallel streams used for forest building + n_streams : int (default = 4) + Number of parallel streams requested for forest building. Distributed + training currently builds trees serially to preserve collective order. workers : optional, list of strings Dask addresses of workers to use for computation. If None, all available Dask workers will be used. random_state : int (default = None) Seed for the random number generator. Unseeded by default. - ignore_empty_partitions: Boolean (default = False) - Specify behavior when a worker does not hold any data - while splitting. When True, it returns the results from workers - with data (the number of trained estimators will be less than - n_estimators) When False, throws a RuntimeError. - """ def __init__( @@ -117,7 +94,6 @@ def __init__( verbose=False, n_estimators=100, random_state=None, - ignore_empty_partitions=False, **kwargs, ): super().__init__(client=client, verbose=verbose, **kwargs) @@ -128,7 +104,6 @@ def __init__( workers=workers, n_estimators=n_estimators, base_seed=random_state, - ignore_empty_partitions=ignore_empty_partitions, **kwargs, ) @@ -138,12 +113,11 @@ def _construct_rf(n_estimators, random_state, **kwargs): n_estimators=n_estimators, random_state=random_state, **kwargs ) - def fit(self, X, y, broadcast_data=False): + def fit(self, X, y): """ Fit the input data with a Random Forest regression model - IMPORTANT: X is expected to be partitioned with at least one partition - on each Dask worker being used by the forest (self.workers). + Only workers holding one or more training rows participate in fitting. When persisting data, you can use `cuml.dask.common.utils.persist_across_workers` to simplify this: @@ -175,16 +149,11 @@ def fit(self, X, y, broadcast_data=False): y : Dask cuDF DataFrame or CuPy backed Dask Array (n_rows, 1) Labels of training examples. **y must be partitioned the same way as X** - broadcast_data : bool, optional (default = False) - When set to True, the whole dataset is broadcasted - to train the workers, otherwise each worker - is trained on its partition """ self.internal_model = None self._fit( model=self.rfs, dataset=(X, y), - broadcast_data=broadcast_data, ) return self @@ -195,7 +164,6 @@ def predict( default_chunk_size=None, align_bytes=None, delayed=True, - broadcast_data=False, ): """ Predicts the regressor outputs for X. @@ -206,7 +174,7 @@ def predict( Distributed dense matrix (floats or doubles) of shape (n_samples, n_features). layout : string (default = 'depth_first') - Specifies the in-memory layout of nodes in FIL forests. Options: + Specifies the in-memory layout of nodes in nvForest models. Options: 'depth_first', 'layered', 'breadth_first'. default_chunk_size : int, optional (default = None) Determines how batches are further subdivided for parallel processing. @@ -220,27 +188,12 @@ def predict( delayed : bool (default = True) Whether to do a lazy prediction (and return Delayed objects) or an eagerly executed one. - broadcast_data : bool (default = False) - If False, the trees are merged in a single model before the workers - perform inference on their share of the prediction workload. - When True, trees aren't merged. Instead each worker infers on the - whole prediction workload using its available trees. The results are - reduced on the client. May be advantageous when the model is larger - than the data used for inference. Returns ------- y : Dask cuDF dataframe or CuPy backed Dask Array (n_rows, 1) """ - if broadcast_data: - return self.partial_inference( - X, - layout=layout, - default_chunk_size=default_chunk_size, - align_bytes=align_bytes, - delayed=delayed, - ) - return self._predict_using_fil( + return self._predict_using_nvforest( X, layout=layout, default_chunk_size=default_chunk_size, @@ -248,20 +201,6 @@ def predict( delayed=delayed, ) - def partial_inference(self, X, delayed, **kwargs): - partial_infs = self._partial_inference( - X=X, op_type="regression", delayed=delayed, **kwargs - ) - workers_weights = self._get_workers_weights() - merged_regressions = dask.array.average( - partial_infs, axis=1, weights=workers_weights - ) - - if delayed: - return merged_regressions - else: - return merged_regressions.persist() - def get_params(self, deep=True): """ Returns the value of all parameters diff --git a/python/cuml/cuml/ensemble/randomforest_common.pyx b/python/cuml/cuml/ensemble/randomforest_common.pyx index 502150b9fc..7876104f94 100644 --- a/python/cuml/cuml/ensemble/randomforest_common.pyx +++ b/python/cuml/cuml/ensemble/randomforest_common.pyx @@ -419,6 +419,9 @@ class BaseRandomForestModel(InteropMixin, Base): ) cdef int n_rows = X.shape[0] cdef int n_cols = X.shape[1] + cdef int parameter_n_rows = getattr( + self, "_distributed_n_rows", n_rows + ) cdef level_enum verbose = self._verbose_level cdef int n_classes = self.n_classes_ if is_classifier else 0 cdef bool input_row_major = not X.flags.f_contiguous @@ -469,19 +472,19 @@ class BaseRandomForestModel(InteropMixin, Base): ) cdef int min_samples_leaf = ( self.min_samples_leaf if isinstance(self.min_samples_leaf, int) - else math.ceil(self.min_samples_leaf * n_rows) + else math.ceil(self.min_samples_leaf * parameter_n_rows) ) cdef int min_samples_split = ( self.min_samples_split if isinstance(self.min_samples_split, int) - else max(2, math.ceil(self.min_samples_split * n_rows)) + else max(2, math.ceil(self.min_samples_split * parameter_n_rows)) ) cdef int n_bins - if self.n_bins > n_rows: + if self.n_bins > parameter_n_rows: warnings.warn("The number of bins, `n_bins` is greater than " "the number of samples used for training. " "Changing `n_bins` to number of training samples.") - n_bins = n_rows + n_bins = parameter_n_rows else: n_bins = self.n_bins @@ -503,7 +506,9 @@ class BaseRandomForestModel(InteropMixin, Base): ) cdef TreeliteModelHandle tl_handle - handle = get_handle(n_streams=n_streams_c) + handle = getattr(self, "_raft_handle", None) + if handle is None: + handle = get_handle(n_streams=n_streams_c) cdef handle_t* handle_ = handle.getHandle() # Store oob_score in C variable for nogil block @@ -604,7 +609,7 @@ class BaseRandomForestModel(InteropMixin, Base): TreeliteFreeModel(tl_handle), "Failed to free Treelite model:" ) - self._n_samples = y.shape[0] + self._n_samples = parameter_n_rows self._n_samples_bootstrap = ( self._n_samples if self.max_samples is None else max(round(self._n_samples * self.max_samples), 1) diff --git a/python/cuml/cuml/ensemble/randomforestclassifier.py b/python/cuml/cuml/ensemble/randomforestclassifier.py index 53ce62ce7c..271d445ff9 100644 --- a/python/cuml/cuml/ensemble/randomforestclassifier.py +++ b/python/cuml/cuml/ensemble/randomforestclassifier.py @@ -259,6 +259,7 @@ def fit(self, X, y, sample_weight=None) -> "RandomForestClassifier": """ Perform Random Forest Classification on the input data """ + classes = getattr(self, "_distributed_classes", True) X, y, sample_weight, classes = check_inputs( self, X, @@ -268,7 +269,7 @@ def fit(self, X, y, sample_weight=None) -> "RandomForestClassifier": order="A", y_dtype="int32", sample_weight_dtype="float64", - return_classes=True, + return_classes=classes, reset=True, ) self.classes_ = classes diff --git a/python/cuml/tests/dask/test_dask_random_forest.py b/python/cuml/tests/dask/test_dask_random_forest.py index 670d2b3532..1c89c72499 100644 --- a/python/cuml/tests/dask/test_dask_random_forest.py +++ b/python/cuml/tests/dask/test_dask_random_forest.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 import json @@ -40,6 +40,10 @@ def _prep_training_data(c, X_train, y_train, partitions_per_worker): return X_train_df, y_train_df +def _get_treelite_bytes(model): + return model._treelite_model_bytes + + @pytest.mark.parametrize("partitions_per_worker", [3]) def test_rf_classification_multi_class(partitions_per_worker, cluster): # Use CUDA_VISIBLE_DEVICES to control the number of workers @@ -65,7 +69,7 @@ def test_rf_classification_multi_class(partitions_per_worker, cluster): ) cu_rf_params = { - "n_estimators": n_workers * 25, + "n_estimators": 25, "max_depth": 16, "n_bins": 256, "random_state": 10, @@ -75,7 +79,7 @@ def test_rf_classification_multi_class(partitions_per_worker, cluster): c, X_train, y_train, partitions_per_worker ) - cuml_mod = cuRFC_mg(**cu_rf_params, ignore_empty_partitions=True) + cuml_mod = cuRFC_mg(**cu_rf_params) cuml_mod.fit(X_train_df, y_train_df) X_test_dask_array = from_array(X_test) cuml_preds_gpu = cuml_mod.predict(X_test_dask_array).compute() @@ -101,7 +105,7 @@ def test_rf_classification_multi_class(partitions_per_worker, cluster): @pytest.mark.parametrize("dtype", [np.float32, np.float64]) @pytest.mark.parametrize("partitions_per_worker", [5]) -def test_rf_regression_dask_fil(partitions_per_worker, dtype, client): +def test_rf_regression_dask_nvforest(partitions_per_worker, dtype, client): n_workers = len(client.scheduler_info(n_workers=-1)["workers"]) # Use CUDA_VISIBLE_DEVICES to control the number of workers @@ -136,7 +140,7 @@ def test_rf_regression_dask_fil(partitions_per_worker, dtype, client): X_cudf_test = cudf.DataFrame(pd.DataFrame(X_test)) X_test_df = dask_cudf.from_cudf(X_cudf_test, npartitions=n_partitions) - cuml_mod = cuRFR_mg(**cu_rf_params, ignore_empty_partitions=True) + cuml_mod = cuRFR_mg(**cu_rf_params) cuml_mod.fit(X_train_df, y_train_df) cuml_mod_predict = cuml_mod.predict(X_test_df) @@ -187,7 +191,7 @@ def test_rf_classification_dask_array(partitions_per_worker, client): @pytest.mark.parametrize("partitions_per_worker", [5]) -def test_rf_classification_dask_fil_predict_proba( +def test_rf_classification_dask_nvforest_predict_proba( partitions_per_worker, client ): n_workers = len(client.scheduler_info(n_workers=-1)["workers"]) @@ -224,16 +228,18 @@ def test_rf_classification_dask_fil_predict_proba( cu_rf_mg = cuRFC_mg(**cu_rf_params) cu_rf_mg.fit(X_train_df, y_train_df) - fil_preds = cu_rf_mg.predict(X_test_df).compute() - fil_preds = fil_preds.to_numpy() - fil_preds_proba = cu_rf_mg.predict_proba(X_test_df).compute() - fil_preds_proba = fil_preds_proba.to_numpy() - np.testing.assert_equal(fil_preds, np.argmax(fil_preds_proba, axis=1)) + nvforest_preds = cu_rf_mg.predict(X_test_df).compute() + nvforest_preds = nvforest_preds.to_numpy() + nvforest_preds_proba = cu_rf_mg.predict_proba(X_test_df).compute() + nvforest_preds_proba = nvforest_preds_proba.to_numpy() + np.testing.assert_equal( + nvforest_preds, np.argmax(nvforest_preds_proba, axis=1) + ) - y_proba = np.zeros(np.shape(fil_preds_proba)) + y_proba = np.zeros(np.shape(nvforest_preds_proba)) y_proba[:, 1] = y_test y_proba[:, 0] = 1.0 - y_test - fil_mse = mean_squared_error(y_proba, fil_preds_proba) + nvforest_mse = mean_squared_error(y_proba, nvforest_preds_proba) sk_model = skrfc( n_estimators=cu_rf_params["n_estimators"], max_depth=cu_rf_params["max_depth"], @@ -245,11 +251,11 @@ def test_rf_classification_dask_fil_predict_proba( # The threshold is required as the test would intermitently # fail with a max difference of 0.029 between the two mse values - assert fil_mse <= sk_mse + 0.029 + assert nvforest_mse <= sk_mse + 0.029 @pytest.mark.parametrize("model_type", ["classification", "regression"]) -def test_rf_concatenation_dask(client, model_type): +def test_rf_distributed_model(client, model_type): n_workers = len(client.scheduler_info(n_workers=-1)["workers"]) X, y = make_classification( @@ -272,49 +278,60 @@ def test_rf_concatenation_dask(client, model_type): cu_rf_mg = cuRFR_mg(**cu_rf_params) cu_rf_mg.fit(X_df, y_df) - res1 = cu_rf_mg.predict(X_df) - res1.compute() - if cu_rf_mg.internal_model: - treelite_bytes = cu_rf_mg.internal_model._treelite_model_bytes - local_tl = treelite.Model.deserialize_bytes(treelite_bytes) - assert local_tl.num_tree == n_estimators + model = cu_rf_mg.get_combined_model() + treelite_bytes = model._treelite_model_bytes + local_tl = treelite.Model.deserialize_bytes(treelite_bytes) + assert local_tl.num_tree == n_estimators + worker_model_bytes = client.gather( + [ + client.submit(_get_treelite_bytes, model, workers=[worker]) + for worker, model in cu_rf_mg.rfs.items() + ] + ) + assert all(data == worker_model_bytes[0] for data in worker_model_bytes) -@pytest.mark.parametrize("ignore_empty_partitions", [True, False]) -def test_single_input_regression(client, ignore_empty_partitions): +def test_rf_classification_uses_global_classes(client): + n_workers = len(client.scheduler_info(n_workers=-1)["workers"]) + if n_workers < 2: + pytest.skip("This test requires at least two workers") + + rows_per_worker = 100 + y = np.repeat(np.arange(n_workers), rows_per_worker).astype(np.int32) + X = np.column_stack((y, np.arange(y.size))).astype(np.float32) + X_dask, y_dask = _prep_training_data(client, X, y, partitions_per_worker=1) + + model = cuRFC_mg( + n_estimators=1, + bootstrap=False, + max_depth=4, + n_bins=max(2, n_workers), + random_state=42, + ).fit(X_dask, y_dask) + + np.testing.assert_array_equal( + model.get_combined_model().classes_, np.arange(n_workers) + ) + + +def test_single_input_regression(client): X, y = make_classification(n_samples=1, n_classes=1) X = X.astype(np.float32) y = y.astype(np.float32) X, y = _prep_training_data(client, X, y, partitions_per_worker=2) - cu_rf_mg = cuRFR_mg( - n_bins=1, - ignore_empty_partitions=ignore_empty_partitions, - ) - - if ( - ignore_empty_partitions - or len(client.scheduler_info(n_workers=-1)["workers"].keys()) == 1 - ): - cu_rf_mg.fit(X, y) - cuml_mod_predict = cu_rf_mg.predict(X) - cuml_mod_predict = cp.asnumpy(cp.array(cuml_mod_predict.compute())) - y = cp.asnumpy(cp.array(y.compute())) - assert y[0] == cuml_mod_predict[0] - - else: - with pytest.raises(ValueError): - cu_rf_mg.fit(X, y) + cu_rf_mg = cuRFR_mg(n_bins=1) + cu_rf_mg.fit(X, y) + cuml_mod_predict = cu_rf_mg.predict(X) + cuml_mod_predict = cp.asnumpy(cp.array(cuml_mod_predict.compute())) + y = cp.asnumpy(cp.array(y.compute())) + assert y[0] == cuml_mod_predict[0] @pytest.mark.parametrize("max_depth", [1, 2, 3, 5, 10, 15, 20]) -@pytest.mark.parametrize("n_estimators", [5, 10, 20]) +@pytest.mark.parametrize("n_estimators", [1, 5, 10, 20]) def test_rf_data_count(client, max_depth, n_estimators): n_workers = len(client.scheduler_info(n_workers=-1)["workers"]) - if n_estimators < n_workers: - err_msg = "n_estimators cannot be lower than number of dask workers" - pytest.xfail(err_msg) - n_samples_per_worker = 350 X, y = make_classification( @@ -355,8 +372,8 @@ def check_count(node, nodes): for tree in json_obj["trees"]: nodes = tree["nodes"] - # The root's count should be equal to the number of rows in the data - assert nodes[0]["data_count"] == n_samples_per_worker + # The root contains rows from the complete distributed dataset. + assert nodes[0]["data_count"] == n_samples_per_worker * n_workers # Check that the data_count accumulates properly as you move up the tree for node in nodes: check_count(node, nodes) @@ -371,7 +388,7 @@ def test_unlimited_max_depth_classifier(client): y = y.astype(np.int32) X_dask, y_dask = _prep_training_data(client, X, y, partitions_per_worker=1) - clf = cuRFC_mg(n_estimators=n_workers * 5, max_depth=None) + clf = cuRFC_mg(n_estimators=5, max_depth=None) clf.fit(X_dask, y_dask) preds = cp.asnumpy(cp.array(clf.predict(X_dask).compute())) assert len(preds) == len(y) @@ -386,22 +403,17 @@ def test_unlimited_max_depth_regressor(client): y = y.astype(np.float32) X_dask, y_dask = _prep_training_data(client, X, y, partitions_per_worker=1) - reg = cuRFR_mg(n_estimators=n_workers * 5, max_depth=None) + reg = cuRFR_mg(n_estimators=5, max_depth=None) reg.fit(X_dask, y_dask) preds = cp.asnumpy(cp.array(reg.predict(X_dask).compute())) assert len(preds) == len(y) @pytest.mark.parametrize("estimator_type", ["regression", "classification"]) -def test_rf_get_combined_model_right_aftter_fit(client, estimator_type): +def test_rf_get_model_right_after_fit(client, estimator_type): max_depth = 3 n_estimators = 5 - n_workers = len(client.scheduler_info(n_workers=-1)["workers"]) - if n_estimators < n_workers: - err_msg = "n_estimators cannot be lower than number of dask workers" - pytest.xfail(err_msg) - X, y = make_classification() X = X.astype(np.float32) if estimator_type == "classification": @@ -437,78 +449,3 @@ def test_rf_get_combined_model_right_aftter_fit(client, estimator_type): assert isinstance(single_gpu_model, cuRFR_sg) else: assert False - - -@pytest.mark.parametrize("model_type", ["classification", "regression"]) -@pytest.mark.parametrize("fit_broadcast", [True, False]) -@pytest.mark.parametrize("transform_broadcast", [True, False]) -def test_rf_broadcast(model_type, fit_broadcast, transform_broadcast, client): - # Use CUDA_VISIBLE_DEVICES to control the number of workers - workers = list(client.scheduler_info(n_workers=-1)["workers"].keys()) - n_workers = len(workers) - - if model_type == "classification": - X, y = make_classification( - n_samples=n_workers * 10000, - n_features=20, - n_informative=15, - n_classes=4, - n_clusters_per_class=1, - random_state=999, - ) - y = y.astype(np.int32) - else: - X, y = make_regression( - n_samples=n_workers * 10000, - n_features=20, - n_informative=5, - random_state=123, - ) - y = y.astype(np.float32) - X = X.astype(np.float32) - - X_train, X_test, y_train, y_test = train_test_split( - X, y, test_size=n_workers * 100, random_state=123 - ) - - X_train_df, y_train_df = _prep_training_data(client, X_train, y_train, 1) - X_test_dask_array = from_array(X_test) - - n_estimators = n_workers * 8 - - if model_type == "classification": - cuml_mod = cuRFC_mg( - n_estimators=n_estimators, - max_depth=8, - n_bins=16, - ignore_empty_partitions=True, - ) - cuml_mod.fit(X_train_df, y_train_df, broadcast_data=fit_broadcast) - cuml_mod_predict = cuml_mod.predict( - X_test_dask_array, broadcast_data=transform_broadcast - ) - - cuml_mod_predict = cuml_mod_predict.compute() - cuml_mod_predict = cp.asnumpy(cuml_mod_predict) - acc_score = accuracy_score(cuml_mod_predict, y_test, normalize=True) - assert acc_score >= 0.68 - - else: - cuml_mod = cuRFR_mg( - n_estimators=n_estimators, - max_depth=8, - n_bins=16, - ignore_empty_partitions=True, - ) - cuml_mod.fit(X_train_df, y_train_df, broadcast_data=fit_broadcast) - cuml_mod_predict = cuml_mod.predict( - X_test_dask_array, broadcast_data=transform_broadcast - ) - - cuml_mod_predict = cuml_mod_predict.compute() - cuml_mod_predict = cp.asnumpy(cuml_mod_predict) - acc_score = r2_score(y_test, cuml_mod_predict) - assert acc_score >= 0.72 - - if transform_broadcast: - assert cuml_mod.internal_model is None