diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index bf3c949ffb..bde93003b7 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -219,6 +219,12 @@ if(BUILD_CUML_MG_TESTS AND NOT SINGLEGPU) set(BUILD_CUML_MPI_COMMS ON) endif() +if(BUILD_CUML_MPI_COMMS) + find_package(MPI COMPONENTS CXX) + find_package(ucx REQUIRED) + find_package(ucxx REQUIRED) +endif() + if(USE_CCACHE) set(CMAKE_C_COMPILER_LAUNCHER ccache) set(CMAKE_CXX_COMPILER_LAUNCHER ccache) diff --git a/cpp/src/decisiontree/batched-levelalgo/builder.cuh b/cpp/src/decisiontree/batched-levelalgo/builder.cuh index e9a7996b65..59749d826c 100644 --- a/cpp/src/decisiontree/batched-levelalgo/builder.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/builder.cuh @@ -15,12 +15,14 @@ #include #include +#include #include #include #include #include +#include #include namespace ML { @@ -36,19 +38,27 @@ class NodeQueue { const DecisionTreeParams params; std::shared_ptr> tree; std::vector node_instances_; + std::vector global_node_counts_; 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 sampled_rows, + int num_outputs, + size_t global_sampled_rows = 0) : params(params), tree(std::make_shared>()) { + if (global_sampled_rows == 0) { global_sampled_rows = sampled_rows; } 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}); + global_node_counts_.reserve(max_nodes); + global_node_counts_.push_back(global_sampled_rows); if (this->IsExpandable(tree->sparsetree.back(), 0)) { work_items_.emplace_back(NodeWorkItem{0, 0, node_instances_.back()}); } @@ -80,49 +90,56 @@ class NodeQueue { } template - void Push(const std::vector& work_items, SplitT* h_splits) + void Push(const std::vector& work_items, SplitT* h_splits, SplitT* h_local_splits) { // Update node queue based on splits for (std::size_t i = 0; i < work_items.size(); i++) { - auto split = h_splits[i]; - auto item = work_items[i]; - auto parent_range = node_instances_.at(item.idx); - if (SplitNotValid( - split, params.min_impurity_decrease, params.min_samples_leaf, parent_range.count)) { + auto global_split = h_splits[i]; + auto item = work_items[i]; + auto parent_range = node_instances_.at(item.idx); + auto parent_global_count = global_node_counts_.at(item.idx); + if (SplitNotValid(global_split, + params.min_impurity_decrease, + params.min_samples_leaf, + parent_global_count)) { continue; } if (params.max_leaves != -1 && tree->leaf_counter >= params.max_leaves) break; - // parent - tree->sparsetree.at(item.idx) = NodeT::CreateSplitNode(split.colid, - split.quesval, - split.best_metric_val, + auto local_split = h_local_splits[i]; + auto left_global_count = std::size_t(global_split.nLeft); + auto right_global_count = parent_global_count - left_global_count; + auto left_local_count = std::size_t(local_split.nLeft); + auto right_local_count = parent_range.count - left_local_count; + + tree->sparsetree.at(item.idx) = NodeT::CreateSplitNode(global_split.colid, + global_split.quesval, + global_split.best_metric_val, int64_t(tree->sparsetree.size()), - parent_range.count); + parent_global_count); tree->leaf_counter++; - // left - tree->sparsetree.emplace_back(NodeT::CreateLeafNode(split.nLeft)); - node_instances_.emplace_back(InstanceRange{parent_range.begin, std::size_t(split.nLeft)}); - // Do not add a work item if this child is definitely a leaf - if (this->IsExpandable(tree->sparsetree.back(), item.depth + 1)) { + tree->sparsetree.emplace_back(NodeT::CreateLeafNode(left_global_count)); + node_instances_.emplace_back(InstanceRange{parent_range.begin, left_local_count}); + global_node_counts_.push_back(left_global_count); + if (item.depth + 1 < params.max_depth && int(left_global_count) >= params.min_samples_split && + (params.max_leaves == -1 || tree->leaf_counter < params.max_leaves)) { work_items_.emplace_back( NodeWorkItem{tree->sparsetree.size() - 1, item.depth + 1, node_instances_.back()}); } - // right - tree->sparsetree.emplace_back(NodeT::CreateLeafNode(parent_range.count - split.nLeft)); + tree->sparsetree.emplace_back(NodeT::CreateLeafNode(right_global_count)); node_instances_.emplace_back( - InstanceRange{parent_range.begin + split.nLeft, parent_range.count - split.nLeft}); - - // Do not add a work item if this child is definitely a leaf - if (this->IsExpandable(tree->sparsetree.back(), item.depth + 1)) { + InstanceRange{parent_range.begin + left_local_count, right_local_count}); + global_node_counts_.push_back(right_global_count); + if (item.depth + 1 < params.max_depth && + int(right_global_count) >= params.min_samples_split && + (params.max_leaves == -1 || tree->leaf_counter < params.max_leaves)) { work_items_.emplace_back( NodeWorkItem{tree->sparsetree.size() - 1, item.depth + 1, node_instances_.back()}); } - // update depth tree->depth_counter = max(tree->depth_counter, item.depth + 1); } } @@ -162,12 +179,12 @@ struct Builder { IdxT* n_nodes; /** buffer of segmented histograms*/ BinT* histograms; - /** threadblock arrival count */ - int* done_count; /** mutex array used for atomically updating best split */ int* mutex; /** best splits for the current batch of nodes */ SplitT* splits; + /** local split counts for distributed partitioning */ + SplitT* local_splits; /** current batch of nodes */ NodeWorkItem* d_work_items; /** device AOS to map CTA blocks along dimx to nodes of a batch */ @@ -178,6 +195,10 @@ struct Builder { int max_blocks_dimx = 0; /** host array of splits */ SplitT* h_splits; + /** host array of local split counts for distributed partitioning */ + SplitT* h_local_splits; + /** packed histogram buffer used by distributed all-reduce */ + void* packed_histograms; /** number of blocks used to parallelize column-wise computations */ int n_blks_for_cols = 10; /** Memory alignment value */ @@ -187,6 +208,10 @@ struct Builder { rmm::device_uvector d_buff; /** pinned host buffer to store the trained nodes */ ML::pinned_host_vector h_buff; + /** true when a communicator with more than one rank is available */ + bool distributed; + /** global root sample count in distributed mode */ + std::size_t global_sampled_rows; Builder(const raft::handle_t& handle, cudaStream_t s, @@ -214,8 +239,23 @@ struct Builder { row_ids->data(), n_classes}, quantiles(q), - d_buff(0, builder_stream) + d_buff(0, builder_stream), + distributed(raft::resource::comms_initialized(handle) && handle.get_comms().get_size() > 1), + global_sampled_rows(row_ids->size()) { + if (distributed) { + auto local_count = static_cast(row_ids->size()); + rmm::device_uvector count_buffer(1, builder_stream); + raft::update_device(count_buffer.data(), &local_count, 1, builder_stream); + handle.get_comms().allreduce( + count_buffer.data(), count_buffer.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 row-count all-reduce."); + auto global_count = std::uint64_t{0}; + raft::update_host(&global_count, count_buffer.data(), 1, builder_stream); + handle.sync_stream(builder_stream); + global_sampled_rows = global_count; + } max_blocks_dimx = 1 + params.max_batch_size + dataset.n_sampled_rows / TPB_DEFAULT; ASSERT(q.quantiles_array != nullptr && q.n_bins_array != nullptr, "Currently quantiles need to be computed before this call!"); @@ -266,13 +306,13 @@ struct Builder { size_t max_len_histograms = max_batch * params.max_n_bins * n_blks_for_cols * dataset.num_outputs; - d_wsize += calculateAlignedBytes(sizeof(IdxT)); // n_nodes - d_wsize += calculateAlignedBytes(sizeof(BinT) * max_len_histograms); // histograms - d_wsize += calculateAlignedBytes(sizeof(int) * max_batch * n_blks_for_cols); // done_count - d_wsize += calculateAlignedBytes(sizeof(int) * max_batch); // mutex - d_wsize += calculateAlignedBytes(sizeof(SplitT) * max_batch); // splits - d_wsize += calculateAlignedBytes(sizeof(NodeWorkItem) * max_batch); // d_work_Items - d_wsize += // workload_info + d_wsize += calculateAlignedBytes(sizeof(IdxT)); // n_nodes + d_wsize += calculateAlignedBytes(sizeof(BinT) * max_len_histograms); // histograms + d_wsize += calculateAlignedBytes(sizeof(int) * max_batch); // mutex + d_wsize += calculateAlignedBytes(sizeof(SplitT) * max_batch); // splits + d_wsize += calculateAlignedBytes(sizeof(SplitT) * max_batch); // local_splits + d_wsize += calculateAlignedBytes(sizeof(NodeWorkItem) * max_batch); // d_work_Items + d_wsize += // workload_info calculateAlignedBytes(sizeof(WorkloadInfo) * max_blocks_dimx); d_wsize += calculateAlignedBytes(sizeof(IdxT) * max_batch * dataset.n_sampled_cols); // colids @@ -280,6 +320,12 @@ struct Builder { h_wsize += // h_workload_info calculateAlignedBytes(sizeof(WorkloadInfo) * max_blocks_dimx); h_wsize += calculateAlignedBytes(sizeof(SplitT) * max_batch); // splits + h_wsize += calculateAlignedBytes(sizeof(SplitT) * max_batch); // local_splits + if constexpr (std::is_same_v) { + d_wsize += calculateAlignedBytes(sizeof(int) * max_len_histograms); + } else { + d_wsize += calculateAlignedBytes(sizeof(double) * max_len_histograms * 2); + } return std::make_pair(d_wsize, h_wsize); } @@ -295,8 +341,7 @@ struct Builder { { raft::common::nvtx::range fun_scope( "Builder::assignWorkspace @builder.cuh [batched-levelalgo]"); - auto max_batch = params.max_batch_size; - auto n_col_blks = n_blks_for_cols; + auto max_batch = params.max_batch_size; size_t max_len_histograms = max_batch * (params.max_n_bins) * n_blks_for_cols * dataset.num_outputs; // device @@ -304,12 +349,12 @@ struct Builder { d_wspace += calculateAlignedBytes(sizeof(IdxT)); histograms = reinterpret_cast(d_wspace); d_wspace += calculateAlignedBytes(sizeof(BinT) * max_len_histograms); - done_count = reinterpret_cast(d_wspace); - d_wspace += calculateAlignedBytes(sizeof(int) * max_batch * n_col_blks); mutex = reinterpret_cast(d_wspace); d_wspace += calculateAlignedBytes(sizeof(int) * max_batch); splits = reinterpret_cast(d_wspace); d_wspace += calculateAlignedBytes(sizeof(SplitT) * max_batch); + local_splits = reinterpret_cast(d_wspace); + d_wspace += calculateAlignedBytes(sizeof(SplitT) * max_batch); d_work_items = reinterpret_cast(d_wspace); d_wspace += calculateAlignedBytes(sizeof(NodeWorkItem) * max_batch); workload_info = reinterpret_cast*>(d_wspace); @@ -317,8 +362,6 @@ struct Builder { colids = reinterpret_cast(d_wspace); d_wspace += calculateAlignedBytes(sizeof(IdxT) * max_batch * dataset.n_sampled_cols); - RAFT_CUDA_TRY( - cudaMemsetAsync(done_count, 0, sizeof(int) * max_batch * n_col_blks, builder_stream)); RAFT_CUDA_TRY(cudaMemsetAsync(mutex, 0, sizeof(int) * max_batch, builder_stream)); // host @@ -326,6 +369,9 @@ struct Builder { h_wspace += calculateAlignedBytes(sizeof(WorkloadInfo) * max_blocks_dimx); h_splits = reinterpret_cast(h_wspace); h_wspace += calculateAlignedBytes(sizeof(SplitT) * max_batch); + h_local_splits = reinterpret_cast(h_wspace); + h_wspace += calculateAlignedBytes(sizeof(SplitT) * max_batch); + packed_histograms = reinterpret_cast(d_wspace); } /** @@ -338,11 +384,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); + params, this->maxNodes(), dataset.n_sampled_rows, dataset.num_outputs, global_sampled_rows); while (queue.HasWork()) { auto work_items = queue.Pop(); auto [splits_host_ptr, splits_count] = doSplit(work_items); - queue.Push(work_items, splits_host_ptr); + queue.Push(work_items, splits_host_ptr, h_local_splits); } auto tree = queue.GetTree(); this->SetLeafPredictions(tree, queue.GetInstanceRanges()); @@ -353,23 +399,19 @@ struct Builder { private: auto updateWorkloadInfo(const std::vector& work_items) { - int n_large_nodes = 0; // large nodes are nodes having training instances larger than block - // size, hence require global memory for histogram construction - int n_blocks_dimx = 0; // gridDim.x required for computeSplitKernel + int n_blocks_dimx = 0; // gridDim.x required for split histogram construction for (std::size_t i = 0; i < work_items.size(); i++) { auto item = work_items[i]; int n_blocks_per_node = std::max(raft::ceildiv(item.instances.count, size_t(TPB_DEFAULT)), size_t(1)); - if (n_blocks_per_node > 1) ++n_large_nodes; - for (int b = 0; b < n_blocks_per_node; b++) { - h_workload_info[n_blocks_dimx + b] = {int(i), n_large_nodes - 1, b, n_blocks_per_node}; + h_workload_info[n_blocks_dimx + b] = {int(i), int(i), b, n_blocks_per_node}; } n_blocks_dimx += n_blocks_per_node; } raft::update_device(workload_info, h_workload_info, n_blocks_dimx, builder_stream); - return std::make_pair(n_blocks_dimx, n_large_nodes); + return n_blocks_dimx; } auto doSplit(const std::vector& work_items) @@ -382,7 +424,7 @@ struct Builder { // get the current set of nodes to be worked upon raft::update_device(d_work_items, work_items.data(), work_items.size(), builder_stream); - auto [n_blocks_dimx, n_large_nodes] = this->updateWorkloadInfo(work_items); + auto n_blocks_dimx = this->updateWorkloadInfo(work_items); // do feature-sampling if (dataset.n_sampled_cols != dataset.N) { @@ -462,7 +504,7 @@ struct Builder { // iterate through a batch of columns (to reduce the memory pressure) and // compute the best split at the end for (IdxT c = 0; c < dataset.n_sampled_cols; c += n_blks_for_cols) { - computeSplit(c, n_blocks_dimx, n_large_nodes); + computeSplit(c, n_blocks_dimx, work_items.size()); RAFT_CUDA_TRY(cudaPeekAtLastError()); } @@ -476,10 +518,12 @@ struct Builder { d_work_items, work_items.size(), splits, + local_splits, builder_stream); RAFT_CUDA_TRY(cudaPeekAtLastError()); raft::common::nvtx::pop_range(); raft::update_host(h_splits, splits, work_items.size(), builder_stream); + raft::update_host(h_local_splits, local_splits, work_items.size(), builder_stream); handle.sync_stream(builder_stream); return std::make_tuple(h_splits, work_items.size()); } @@ -488,11 +532,9 @@ struct Builder { { size_t smem_size_1 = params.max_n_bins * dataset.num_outputs * sizeof(BinT) + // shared_histogram size - params.max_n_bins * sizeof(DataT) + // shared_quantiles size - sizeof(int); // shared_done size - // Extra room for alignment (see alignPointer in - // computeSplitKernel) - smem_size_1 += sizeof(DataT) + 3 * sizeof(int); + params.max_n_bins * sizeof(DataT); // shared_quantiles size + // Extra room for alignment (see alignPointer in the split kernels). + smem_size_1 += sizeof(DataT); // Calculate the shared memory needed for evalBestSplit size_t smem_size_2 = raft::ceildiv(TPB_DEFAULT, raft::WarpSize) * sizeof(SplitT); // Pick the max of two @@ -502,7 +544,31 @@ struct Builder { return smem_size; } - void computeSplit(IdxT col, size_t n_blocks_dimx, size_t n_large_nodes) + void allReduceHistograms(BinT* histograms_to_reduce, std::size_t len_histograms) + { + auto const& comm = handle.get_comms(); + if constexpr (std::is_same_v) { + auto* packed = reinterpret_cast(packed_histograms); + packHistograms(histograms_to_reduce, packed, len_histograms, builder_stream); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + comm.allreduce(packed, packed, len_histograms, raft::comms::op_t::SUM, builder_stream); + ASSERT(comm.sync_stream(builder_stream) == raft::comms::status_t::SUCCESS, + "An error occurred in the distributed RF histogram all-reduce."); + unpackHistograms(packed, histograms_to_reduce, len_histograms, builder_stream); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } else { + auto* packed = reinterpret_cast(packed_histograms); + packHistograms(histograms_to_reduce, packed, len_histograms, builder_stream); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + comm.allreduce(packed, packed, 2 * len_histograms, raft::comms::op_t::SUM, builder_stream); + ASSERT(comm.sync_stream(builder_stream) == raft::comms::status_t::SUCCESS, + "An error occurred in the distributed RF histogram all-reduce."); + unpackHistograms(packed, histograms_to_reduce, len_histograms, builder_stream); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + } + + void computeSplit(IdxT col, size_t n_blocks_dimx, size_t work_items_size) { // if no instances to split, return if (n_blocks_dimx == 0) return; @@ -516,31 +582,42 @@ struct Builder { dim3 grid(n_blocks_dimx, n_blocks_dimy, 1); // required total length (in bins) of the global segmented histograms over all // classes, features and (large)nodes. - int len_histograms = n_bins * n_classes * n_blocks_dimy * n_large_nodes; + int len_histograms = n_bins * n_classes * n_blocks_dimy * work_items_size; RAFT_CUDA_TRY(cudaMemsetAsync(histograms, 0, sizeof(BinT) * len_histograms, builder_stream)); // create the objective function object ObjectiveT objective(dataset.num_outputs, params.min_samples_leaf); - // call the computeSplitKernel - raft::common::nvtx::range kernel_scope("computeSplitKernel @builder.cuh [batched-levelalgo]"); - launchComputeSplitKernel(histograms, - params.max_n_bins, - params.min_samples_split, - params.max_leaves, - dataset, - quantiles, - d_work_items, - col, - colids, - done_count, - mutex, - splits, - objective, - treeid, - workload_info, - seed, - grid, - smem_size, - builder_stream); + raft::common::nvtx::range kernel_scope("split kernels @builder.cuh [batched-levelalgo]"); + launchComputeSplitHistogramKernel(histograms, + params.max_n_bins, + dataset, + quantiles, + d_work_items, + col, + colids, + objective, + treeid, + workload_info, + seed, + grid, + smem_size, + builder_stream); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + if (distributed) { allReduceHistograms(histograms, len_histograms); } + dim3 eval_grid(work_items_size, n_blocks_dimy, 1); + launchEvaluateSplitKernel(histograms, + params.max_n_bins, + dataset, + quantiles, + d_work_items, + col, + colids, + mutex, + splits, + objective, + treeid, + eval_grid, + smem_size, + builder_stream); } // Set the leaf value predictions in batch @@ -554,6 +631,8 @@ struct Builder { std::size_t max_batch_size = min(std::size_t(100000), tree->sparsetree.size()); 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_leaf_histograms(max_batch_size * dataset.num_outputs, + builder_stream); rmm::device_uvector d_leaves(max_batch_size * dataset.num_outputs, builder_stream); ObjectiveT objective(dataset.num_outputs, params.min_samples_leaf); @@ -566,17 +645,31 @@ struct Builder { raft::update_device( d_instance_ranges.data(), instance_ranges.data() + batch_begin, batch_size, builder_stream); + RAFT_CUDA_TRY(cudaMemsetAsync( + d_leaf_histograms.data(), 0, sizeof(BinT) * d_leaf_histograms.size(), builder_stream)); RAFT_CUDA_TRY( cudaMemsetAsync(d_leaves.data(), 0, sizeof(DataT) * d_leaves.size(), builder_stream)); size_t smem_size = sizeof(BinT) * dataset.num_outputs; - launchLeafKernel(objective, - dataset, - d_tree.data(), - d_instance_ranges.data(), - d_leaves.data(), - batch_size, - smem_size, - builder_stream); + launchLeafHistogramKernel(objective, + dataset, + d_tree.data(), + d_instance_ranges.data(), + d_leaf_histograms.data(), + batch_size, + smem_size, + builder_stream); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + if (distributed) { + allReduceHistograms(d_leaf_histograms.data(), batch_size * dataset.num_outputs); + } + launchFinalizeLeafKernel(objective, + d_tree.data(), + d_leaf_histograms.data(), + d_leaves.data(), + batch_size, + dataset.num_outputs, + builder_stream); + RAFT_CUDA_TRY(cudaPeekAtLastError()); raft::update_host(tree->vector_leaf.data() + batch_begin * dataset.num_outputs, d_leaves.data(), batch_size * dataset.num_outputs, diff --git a/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh b/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh index 6f56228d35..f0cd6a47b2 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh @@ -17,6 +17,8 @@ #include +#include + namespace ML { namespace DT { @@ -39,9 +41,8 @@ struct NodeWorkItem { */ template struct WorkloadInfo { - IdxT nodeid; // Node in the batch on which the threadblock needs to work - IdxT large_nodeid; // counts only large nodes (nodes that require more than one block along x-dim - // for histogram calculation) + IdxT nodeid; // Node in the batch on which the threadblock needs to work + IdxT large_nodeid; // legacy field; histogram offsets now use nodeid for all nodes IdxT offset_blockid; // Offset threadblock id among all the blocks that are // working on this node IdxT num_blocks; // Total number of blocks that are working on the node @@ -73,17 +74,27 @@ void launchNodeSplitKernel(const IdxT min_samples_leaf, const NodeWorkItem* work_items, const size_t work_items_size, const Split* splits, + Split* local_splits, 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 launchLeafHistogramKernel(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(ObjectiveT objective, + const NodeT* tree, + const typename ObjectiveT::BinT* leaf_histograms, + DataT* leaves, + int batch_size, + int num_outputs, + cudaStream_t builder_stream); // Returns the lowest index in `array` whose value is greater or equal to `element`. // Values outside the quantile range are clamped to the edge bins: values below the // first quantile return 0, and values above the last quantile return len - 1. @@ -371,25 +382,91 @@ template -void launchComputeSplitKernel(BinT* histograms, - IdxT n_bins, - IdxT min_samples_split, - IdxT max_leaves, - const Dataset& dataset, - const Quantiles& quantiles, - const NodeWorkItem* work_items, - IdxT colStart, - const IdxT* colids, - int* done_count, - int* mutex, - volatile Split* splits, - ObjectiveT& objective, - IdxT treeid, - const WorkloadInfo* workload_info, - uint64_t seed, - dim3 grid, - size_t smem_size, - cudaStream_t builder_stream); +void launchComputeSplitHistogramKernel(BinT* histograms, + IdxT max_n_bins, + const Dataset& dataset, + const Quantiles& quantiles, + const NodeWorkItem* work_items, + IdxT colStart, + const IdxT* colids, + ObjectiveT& objective, + IdxT treeid, + const WorkloadInfo* workload_info, + uint64_t seed, + dim3 grid, + size_t smem_size, + cudaStream_t builder_stream); + +template +void launchEvaluateSplitKernel(BinT* histograms, + IdxT max_n_bins, + const Dataset& dataset, + const Quantiles& quantiles, + const NodeWorkItem* work_items, + IdxT colStart, + const IdxT* colids, + int* mutex, + volatile Split* splits, + ObjectiveT& objective, + IdxT treeid, + dim3 grid, + size_t smem_size, + cudaStream_t builder_stream); + +template +static __global__ void transformHistogramKernel(const InT* in, OutT* out, std::size_t len, OpT op) +{ + std::size_t tid = blockIdx.x * blockDim.x + threadIdx.x; + for (std::size_t i = tid; i < len; i += std::size_t(blockDim.x) * gridDim.x) { + op(in, out, i); + } +} + +inline std::size_t histogramTransformBlocks(std::size_t len) +{ + return std::max(std::size_t{1}, raft::ceildiv(len, 256)); +} + +inline void packHistograms(const CountBin* in, int* out, std::size_t len, cudaStream_t stream) +{ + auto op = [] __device__(const CountBin* in, int* out, std::size_t i) { out[i] = in[i].x; }; + transformHistogramKernel<<>>(in, out, len, op); +} + +inline void unpackHistograms(const int* in, CountBin* out, std::size_t len, cudaStream_t stream) +{ + auto op = [] __device__(const int* in, CountBin* out, std::size_t i) { out[i].x = in[i]; }; + transformHistogramKernel<<>>(in, out, len, op); +} + +inline void packHistograms(const AggregateBin* in, + double* out, + std::size_t len, + cudaStream_t stream) +{ + auto op = [] __device__(const AggregateBin* in, double* out, std::size_t i) { + out[2 * i] = in[i].label_sum; + out[2 * i + 1] = static_cast(in[i].count); + }; + transformHistogramKernel<<>>(in, out, len, op); +} + +inline void unpackHistograms(const double* in, + AggregateBin* out, + std::size_t len, + cudaStream_t stream) +{ + auto op = [] __device__(const double* in, AggregateBin* out, std::size_t i) { + out[i].label_sum = in[2 * i]; + out[i].count = static_cast(in[2 * i + 1]); + }; + transformHistogramKernel<<>>(in, out, len, op); +} } // namespace DT } // namespace ML 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 59677b6caf..97a49ba5cd 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels_impl.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels_impl.cuh @@ -14,6 +14,7 @@ #include #include +#include #include namespace ML { @@ -84,16 +85,31 @@ static __global__ void nodeSplitKernel(const IdxT min_samples_leaf, const DataT min_impurity_decrease, const Dataset dataset, const NodeWorkItem* work_items, - const Split* splits) + const Split* splits, + Split* local_splits) { extern __shared__ char smem[]; const auto work_item = work_items[blockIdx.x]; - const auto split = splits[blockIdx.x]; - if (SplitNotValid( - split, min_impurity_decrease, min_samples_leaf, IdxT(work_item.instances.count))) { - return; + auto split = splits[blockIdx.x]; + if (split.best_metric_val <= min_impurity_decrease) { return; } + + auto* left_count = reinterpret_cast(smem); + if (threadIdx.x == 0) { *left_count = IdxT{0}; } + __syncthreads(); + + auto* col = dataset.data + split.colid * std::size_t(dataset.M); + for (auto i = work_item.instances.begin + threadIdx.x; + i < work_item.instances.begin + work_item.instances.count; + i += blockDim.x) { + auto row = dataset.row_ids[i]; + if (col[row] <= split.quesval) { atomicAdd(left_count, IdxT{1}); } } - partitionSamples(dataset, split, work_item, (char*)smem); + __syncthreads(); + + split.nLeft = *left_count; + local_splits[blockIdx.x] = split; + auto* partition_smem = alignPointer(left_count + 1); + partitionSamples(dataset, split, work_item, partition_smem); } template @@ -105,9 +121,10 @@ void launchNodeSplitKernel(const IdxT min_samples_leaf, const NodeWorkItem* work_items, const size_t work_items_size, const Split* splits, + Split* local_splits, cudaStream_t builder_stream) { - auto constexpr smem_size = 2 * sizeof(IdxT) * TPB; + auto constexpr smem_size = sizeof(IdxT) + 2 * sizeof(IdxT) * TPB + sizeof(IdxT); nodeSplitKernel <<>>(min_samples_leaf, min_samples_split, @@ -115,15 +132,16 @@ void launchNodeSplitKernel(const IdxT min_samples_leaf, min_impurity_decrease, dataset, work_items, - splits); + splits, + local_splits); } -template +template static __global__ void leafKernel(ObjectiveT objective, DatasetT dataset, const NodeT* tree, const InstanceRange* instance_ranges, - DataT* leaves) + typename ObjectiveT::BinT* leaf_histograms) { using BinT = typename ObjectiveT::BinT; extern __shared__ char shared_memory[]; @@ -143,24 +161,59 @@ static __global__ void leafKernel(ObjectiveT objective, } __syncthreads(); if (tid == 0) { - ObjectiveT::SetLeafVector( - histogram, dataset.num_outputs, leaves + dataset.num_outputs * node_id); + auto leaf_histogram = leaf_histograms + dataset.num_outputs * node_id; + for (int i = 0; i < dataset.num_outputs; ++i) { + leaf_histogram[i] = histogram[i]; + } + } +} + +template +static __global__ void finalizeLeafKernel(ObjectiveT objective, + const NodeT* tree, + const typename ObjectiveT::BinT* leaf_histograms, + DataT* leaves, + int num_outputs) +{ + auto node_id = blockIdx.x; + auto leaf = leaves + num_outputs * node_id; + auto& node = tree[node_id]; + if (!node.IsLeaf()) { + for (int i = 0; i < num_outputs; ++i) { + leaf[i] = DataT(0); + } + return; } + auto leaf_histogram = leaf_histograms + num_outputs * node_id; + ObjectiveT::SetLeafVector(leaf_histogram, num_outputs, leaf); } -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 launchFinalizeLeafKernel(ObjectiveT objective, + const NodeT* tree, + const typename ObjectiveT::BinT* leaf_histograms, + DataT* leaves, + int batch_size, + int num_outputs, + cudaStream_t builder_stream) +{ + finalizeLeafKernel<<>>( + objective, tree, leaf_histograms, leaves, num_outputs); +} + +template +void launchLeafHistogramKernel(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) { int num_blocks = batch_size; leafKernel<<>>( - objective, dataset, tree, instance_ranges, leaves); + objective, dataset, tree, instance_ranges, leaf_histograms); } /** @@ -193,44 +246,38 @@ DI BinT pdf_to_cdf(BinT* shared_histogram, IdxT n_bins) return total_aggregate; } +DI int bin_count(CountBin const& bin) { return bin.x; } + +DI int bin_count(AggregateBin const& bin) { return bin.count; } + template -static __global__ void computeSplitKernel(BinT* histograms, - IdxT max_n_bins, - IdxT min_samples_split, - IdxT max_leaves, - const Dataset dataset, - const Quantiles quantiles, - const NodeWorkItem* work_items, - IdxT colStart, - const IdxT* colids, - int* done_count, - int* mutex, - volatile Split* splits, - ObjectiveT objective, - IdxT treeid, - const WorkloadInfo* workload_info, - uint64_t seed) +static __global__ void computeSplitHistogramKernel(BinT* histograms, + IdxT max_n_bins, + const Dataset dataset, + const Quantiles quantiles, + const NodeWorkItem* work_items, + IdxT colStart, + const IdxT* colids, + ObjectiveT objective, + IdxT treeid, + const WorkloadInfo* workload_info, + uint64_t seed) { - // dynamic shared memory extern __shared__ char smem[]; - // Read workload info for this block - WorkloadInfo workload_info_cta = workload_info[blockIdx.x]; - IdxT nid = workload_info_cta.nodeid; - IdxT large_nid = workload_info_cta.large_nodeid; - const auto work_item = work_items[nid]; - auto range_start = work_item.instances.begin; - auto range_len = work_item.instances.count; - - IdxT offset_blockid = workload_info_cta.offset_blockid; - IdxT num_blocks = workload_info_cta.num_blocks; + auto workload_info_cta = workload_info[blockIdx.x]; + IdxT nid = workload_info_cta.nodeid; + const auto work_item = work_items[nid]; + auto range_start = work_item.instances.begin; + auto range_len = work_item.instances.count; + IdxT offset_blockid = workload_info_cta.offset_blockid; + IdxT num_blocks = workload_info_cta.num_blocks; - // obtaining the feature to test split on IdxT col; if (dataset.n_sampled_cols == dataset.N) { col = colStart + blockIdx.y; @@ -239,89 +286,90 @@ static __global__ void computeSplitKernel(BinT* histograms, col = colids[nid * dataset.n_sampled_cols + colIndex]; } - // getting the n_bins for that feature - int n_bins = quantiles.n_bins_array[col]; - - auto end = range_start + range_len; + int n_bins = quantiles.n_bins_array[col]; auto shared_histogram_len = n_bins * objective.NumClasses(); auto* shared_histogram = alignPointer(smem); auto* shared_quantiles = alignPointer(shared_histogram + shared_histogram_len); - auto* shared_done = alignPointer(shared_quantiles + n_bins); IdxT stride = blockDim.x * num_blocks; IdxT tid = threadIdx.x + offset_blockid * blockDim.x; + auto histograms_offset = ((nid * gridDim.y) + blockIdx.y) * max_n_bins * objective.NumClasses(); - // populating shared memory with initial values - for (IdxT i = threadIdx.x; i < shared_histogram_len; i += blockDim.x) + for (IdxT i = threadIdx.x; i < shared_histogram_len; i += blockDim.x) { shared_histogram[i] = BinT(); - for (IdxT b = threadIdx.x; b < n_bins; b += blockDim.x) + } + for (IdxT b = threadIdx.x; b < n_bins; b += blockDim.x) { shared_quantiles[b] = quantiles.quantiles_array[max_n_bins * col + b]; - - // synchronizing above changes across block + } __syncthreads(); - // compute pdf shared histogram for all bins for all classes in shared mem - - // Must be 64 bit - can easily grow larger than a 32 bit int std::size_t col_offset = std::size_t(col) * dataset.M; - for (auto i = range_start + tid; i < end; i += stride) { - // each thread works over a data point and strides to the next + for (auto i = range_start + tid; i < range_start + range_len; i += stride) { auto row = dataset.row_ids[i]; auto data = dataset.data[row + col_offset]; auto label = dataset.labels[row]; - - // `start` is lowest index such that data <= shared_quantiles[start] IdxT start = lower_bound(shared_quantiles, n_bins, data); - // ++shared_histogram[start] BinT::IncrementHistogram(shared_histogram, n_bins, start, label); } - - // synchronizing above changes across block __syncthreads(); - if (num_blocks > 1) { - // update the corresponding global location - auto histograms_offset = - ((large_nid * gridDim.y) + blockIdx.y) * max_n_bins * objective.NumClasses(); - for (IdxT i = threadIdx.x; i < shared_histogram_len; i += blockDim.x) { - BinT::AtomicAdd(histograms + histograms_offset + i, shared_histogram[i]); - } - __threadfence(); // for commit guarantee - __syncthreads(); + for (IdxT i = threadIdx.x; i < shared_histogram_len; i += blockDim.x) { + BinT::AtomicAdd(histograms + histograms_offset + i, shared_histogram[i]); + } +} - // last threadblock will go ahead and compute the best split - bool last = MLCommon::signalDone( - done_count + nid * gridDim.y + blockIdx.y, num_blocks, offset_blockid == 0, shared_done); - // if not the last threadblock, exit - if (!last) return; +template +static __global__ void evaluateSplitKernel(BinT* histograms, + IdxT max_n_bins, + const Dataset dataset, + const Quantiles quantiles, + const NodeWorkItem* work_items, + IdxT colStart, + const IdxT* colids, + int* mutex, + volatile Split* splits, + ObjectiveT objective, + IdxT treeid) +{ + extern __shared__ char smem[]; - // store the complete global histogram in shared memory of last block - for (IdxT i = threadIdx.x; i < shared_histogram_len; i += blockDim.x) - shared_histogram[i] = histograms[histograms_offset + i]; + IdxT nid = blockIdx.x; + IdxT col; + if (dataset.n_sampled_cols == dataset.N) { + col = colStart + blockIdx.y; + } else { + IdxT colIndex = colStart + blockIdx.y; + col = colids[nid * dataset.n_sampled_cols + colIndex]; + } - __syncthreads(); + int n_bins = quantiles.n_bins_array[col]; + auto shared_histogram_len = n_bins * objective.NumClasses(); + auto* shared_histogram = alignPointer(smem); + auto* shared_quantiles = alignPointer(shared_histogram + shared_histogram_len); + auto histograms_offset = ((nid * gridDim.y) + blockIdx.y) * max_n_bins * objective.NumClasses(); + + for (IdxT i = threadIdx.x; i < shared_histogram_len; i += blockDim.x) { + shared_histogram[i] = histograms[histograms_offset + i]; } + for (IdxT b = threadIdx.x; b < n_bins; b += blockDim.x) { + shared_quantiles[b] = quantiles.quantiles_array[max_n_bins * col + b]; + } + __syncthreads(); - // PDF to CDF inplace in `shared_histogram` + IdxT split_len = 0; for (IdxT c = 0; c < objective.NumClasses(); ++c) { - // left to right scan operation for scanning - // "lesser-than-or-equal" counts - BinT total_sum = pdf_to_cdf(shared_histogram + n_bins * c, n_bins); - // now, `shared_histogram[n_bins * c + i]` will have count of datapoints of class `c` - // that are less than or equal to `shared_quantiles[i]`. + pdf_to_cdf(shared_histogram + n_bins * c, n_bins); + split_len += bin_count(shared_histogram[n_bins * c + n_bins - 1]); } - __syncthreads(); - // calculate the best candidate bins (one for each thread in the block) in current feature and - // corresponding information gain for splitting Split sp = - objective.Gain(shared_histogram, shared_quantiles, col, range_len, n_bins); - + objective.Gain(shared_histogram, shared_quantiles, col, split_len, n_bins); __syncthreads(); - - // calculate best bins among candidate bins per feature using warp reduce - // then atomically update across features to get best split per node - // (in split[nid]) sp.evalBestSplit(smem, splits + nid, mutex + nid); } @@ -331,45 +379,70 @@ template -void launchComputeSplitKernel(BinT* histograms, - IdxT max_n_bins, - IdxT min_samples_split, - IdxT max_leaves, - const Dataset& dataset, - const Quantiles& quantiles, - const NodeWorkItem* work_items, - IdxT colStart, - const IdxT* colids, - int* done_count, - int* mutex, - volatile Split* splits, - ObjectiveT& objective, - IdxT treeid, - const WorkloadInfo* workload_info, - uint64_t seed, - dim3 grid, - size_t smem_size, - cudaStream_t builder_stream) +void launchComputeSplitHistogramKernel(BinT* histograms, + IdxT max_n_bins, + const Dataset& dataset, + const Quantiles& quantiles, + const NodeWorkItem* work_items, + IdxT colStart, + const IdxT* colids, + ObjectiveT& objective, + IdxT treeid, + const WorkloadInfo* workload_info, + uint64_t seed, + dim3 grid, + size_t smem_size, + cudaStream_t builder_stream) { - computeSplitKernel + computeSplitHistogramKernel <<>>(histograms, max_n_bins, - min_samples_split, - max_leaves, dataset, quantiles, work_items, colStart, colids, - done_count, - mutex, - splits, objective, treeid, workload_info, seed); } +template +void launchEvaluateSplitKernel(BinT* histograms, + IdxT max_n_bins, + const Dataset& dataset, + const Quantiles& quantiles, + const NodeWorkItem* work_items, + IdxT colStart, + const IdxT* colids, + int* mutex, + volatile Split* splits, + ObjectiveT& objective, + IdxT treeid, + dim3 grid, + size_t smem_size, + cudaStream_t builder_stream) +{ + evaluateSplitKernel + <<>>(histograms, + max_n_bins, + dataset, + quantiles, + work_items, + colStart, + colids, + mutex, + splits, + objective, + treeid); +} + template void launchNodeSplitKernel<_DataT, _LabelT, _IdxT, TPB_DEFAULT>( const _IdxT min_samples_leaf, const _IdxT min_samples_split, @@ -379,35 +452,57 @@ template void launchNodeSplitKernel<_DataT, _LabelT, _IdxT, TPB_DEFAULT>( const NodeWorkItem* work_items, const size_t work_items_size, const Split<_DataT, _IdxT>* splits, + Split<_DataT, _IdxT>* local_splits, cudaStream_t builder_stream); -template void launchLeafKernel<_DatasetT, _NodeT, _ObjectiveT, _DataT>( +template void launchLeafHistogramKernel<_DatasetT, _NodeT, _ObjectiveT>( _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<_NodeT, _ObjectiveT, _DataT>( + _ObjectiveT objective, + const _NodeT* tree, + const typename _ObjectiveT::BinT* leaf_histograms, _DataT* leaves, int batch_size, + int num_outputs, + cudaStream_t builder_stream); + +template void +launchComputeSplitHistogramKernel<_DataT, _LabelT, _IdxT, TPB_DEFAULT, _ObjectiveT, _BinT>( + _BinT* histograms, + _IdxT max_n_bins, + const Dataset<_DataT, _LabelT, _IdxT>& dataset, + const Quantiles<_DataT, _IdxT>& quantiles, + const NodeWorkItem* work_items, + _IdxT colStart, + const _IdxT* colids, + _ObjectiveT& objective, + _IdxT treeid, + const WorkloadInfo<_IdxT>* workload_info, + uint64_t seed, + dim3 grid, size_t smem_size, cudaStream_t builder_stream); -template void launchComputeSplitKernel<_DataT, _LabelT, _IdxT, TPB_DEFAULT, _ObjectiveT, _BinT>( +template void launchEvaluateSplitKernel<_DataT, _LabelT, _IdxT, TPB_DEFAULT, _ObjectiveT, _BinT>( _BinT* histograms, - _IdxT n_bins, - _IdxT min_samples_split, - _IdxT max_leaves, + _IdxT max_n_bins, const Dataset<_DataT, _LabelT, _IdxT>& dataset, const Quantiles<_DataT, _IdxT>& quantiles, const NodeWorkItem* work_items, _IdxT colStart, const _IdxT* colids, - int* done_count, int* mutex, volatile Split<_DataT, _IdxT>* splits, _ObjectiveT& objective, _IdxT treeid, - const WorkloadInfo<_IdxT>* workload_info, - uint64_t seed, dim3 grid, size_t smem_size, cudaStream_t builder_stream); diff --git a/cpp/src/decisiontree/batched-levelalgo/quantiles.cuh b/cpp/src/decisiontree/batched-levelalgo/quantiles.cuh index 362c4004be..203eb62d27 100644 --- a/cpp/src/decisiontree/batched-levelalgo/quantiles.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/quantiles.cuh @@ -13,6 +13,7 @@ #include #include #include +#include #include #include @@ -24,14 +25,31 @@ #include #include +#include #include #include +#include +#include namespace ML { namespace DT { namespace detail { +inline std::vector proportionalSampleCounts( + std::vector const& row_counts, std::uint64_t global_rows, int sample_count) +{ + std::vector result(row_counts.size()); + std::uint64_t prefix = 0; + for (std::size_t i = 0; i < row_counts.size(); ++i) { + auto begin = (static_cast(sample_count) * prefix) / global_rows; + prefix += row_counts[i]; + auto end = (static_cast(sample_count) * prefix) / global_rows; + result[i] = static_cast(end - begin); + } + return result; +} + template static __global__ void gatherUniformSampledColumnKernel( T* out, const T* data, int sample_count, int n_rows, int col, uint64_t seed) @@ -124,17 +142,49 @@ CUML_EXPORT QuantileReturnValue computeQuantiles(const raft::handle_t& handle raft::common::nvtx::push_range("computeQuantiles"); RAFT_EXPECTS(data != nullptr, "data pointer must not be null"); RAFT_EXPECTS(max_n_bins > 0, "max_n_bins must be positive"); - RAFT_EXPECTS(n_rows > 0, "n_rows must be positive"); + RAFT_EXPECTS(n_rows >= 0, "n_rows must be non-negative"); 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(); - int64_t size = static_cast(max_n_bins) * oversampling_factor; - int sample_count = - static_cast(std::min(static_cast(n_rows), std::max(1, size))); + 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; + int64_t size = static_cast(max_n_bins) * oversampling_factor; + auto target_sample = std::max(1, size); + + std::uint64_t global_rows = static_cast(n_rows); + std::vector rank_rows(1, global_rows); + std::vector rank_sample_counts(1, 0); + std::vector rank_sample_displs(1, 0); + + if (distributed) { + rmm::device_uvector row_counts(comm_size, stream); + auto local_rows = static_cast(n_rows); + raft::update_device(row_counts.data(), &local_rows, 1, stream); + handle.get_comms().allgather(row_counts.data(), row_counts.data(), 1, stream); + ASSERT(handle.get_comms().sync_stream(stream) == raft::comms::status_t::SUCCESS, + "An error occurred in the distributed RF quantile row-count all-gather."); + rank_rows.resize(comm_size); + raft::update_host(rank_rows.data(), row_counts.data(), comm_size, stream); + handle.sync_stream(stream); + global_rows = std::accumulate(rank_rows.begin(), rank_rows.end(), std::uint64_t{0}); + } + RAFT_EXPECTS(global_rows > 0, "global row count must be positive"); + + int sample_count = static_cast( + std::min(global_rows, static_cast(target_sample))); + rank_sample_counts = detail::proportionalSampleCounts(rank_rows, global_rows, sample_count); + rank_sample_displs.resize(comm_size); + for (int i = 1; i < comm_size; ++i) { + rank_sample_displs[i] = rank_sample_displs[i - 1] + rank_sample_counts[i - 1]; + } + int local_sample_count = static_cast(rank_sample_counts[rank]); rmm::device_uvector sampled_column(sample_count, stream); rmm::device_uvector sorted_sample(sample_count, stream); + rmm::device_uvector local_sampled_column( + distributed ? std::max(1, local_sample_count) : sample_count, stream); auto quantiles_array = std::make_shared>(n_cols * max_n_bins, stream); auto n_bins_array = std::make_shared>(n_cols, stream); @@ -150,28 +200,50 @@ CUML_EXPORT QuantileReturnValue computeQuantiles(const raft::handle_t& handle rmm::device_uvector d_temp_storage(temp_storage_bytes, stream); int n_threads = 256; - int n_blocks = raft::ceildiv(sample_count, n_threads); + int n_blocks = raft::ceildiv(std::max(1, local_sample_count), n_threads); n_blocks = std::min(n_blocks, 1024); for (int col = 0; col < n_cols; col++) { raft::common::nvtx::push_range("sample quantile column"); - if (sample_count == n_rows) { - RAFT_CUDA_TRY(cudaMemcpyAsync(sampled_column.data(), - data + static_cast(col) * n_rows, - sizeof(T) * n_rows, - cudaMemcpyDeviceToDevice, - stream)); + T* sort_input = sampled_column.data(); + if (distributed) { + if (local_sample_count > 0 && local_sample_count == n_rows) { + RAFT_CUDA_TRY(cudaMemcpyAsync(local_sampled_column.data(), + data + static_cast(col) * n_rows, + sizeof(T) * n_rows, + cudaMemcpyDeviceToDevice, + stream)); + } else if (local_sample_count > 0) { + detail::gatherUniformSampledColumnKernel<<>>( + local_sampled_column.data(), data, local_sample_count, n_rows, col, seed); + RAFT_CUDA_TRY(cudaGetLastError()); + } + handle.get_comms().allgatherv(local_sampled_column.data(), + sampled_column.data(), + rank_sample_counts.data(), + rank_sample_displs.data(), + stream); + ASSERT(handle.get_comms().sync_stream(stream) == raft::comms::status_t::SUCCESS, + "An error occurred in the distributed RF quantile sample all-gather."); } else { - detail::gatherUniformSampledColumnKernel<<>>( - sampled_column.data(), data, sample_count, n_rows, col, seed); - RAFT_CUDA_TRY(cudaGetLastError()); + if (sample_count == n_rows) { + RAFT_CUDA_TRY(cudaMemcpyAsync(sampled_column.data(), + data + static_cast(col) * n_rows, + sizeof(T) * n_rows, + cudaMemcpyDeviceToDevice, + stream)); + } else { + detail::gatherUniformSampledColumnKernel<<>>( + sampled_column.data(), data, sample_count, n_rows, col, seed); + RAFT_CUDA_TRY(cudaGetLastError()); + } } raft::common::nvtx::pop_range(); raft::common::nvtx::push_range("sort sampled quantile column"); RAFT_CUDA_TRY(cub::DeviceRadixSort::SortKeys((void*)(d_temp_storage.data()), temp_storage_bytes, - sampled_column.data(), + sort_input, sorted_sample.data(), sample_count, 0, diff --git a/cpp/src/randomforest/randomforest.cuh b/cpp/src/randomforest/randomforest.cuh index f0115402aa..eaaf0cb5d8 100644 --- a/cpp/src/randomforest/randomforest.cuh +++ b/cpp/src/randomforest/randomforest.cuh @@ -9,6 +9,7 @@ #include #include +#include #include #include #include @@ -49,6 +50,7 @@ class RandomForest { const cudaStream_t stream) { raft::common::nvtx::range fun_scope("bootstrapping row IDs @randomforest.cuh"); + if (selected_rows->size() == 0) { return; } // Hash these together so they are uncorrelated auto rs = DT::fnv1a32_basis; @@ -65,14 +67,21 @@ class RandomForest { } } - 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); @@ -120,9 +129,13 @@ class RandomForest { bool* bootstrap_masks = nullptr) { 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; - int n_sampled_rows = 0; + bool distributed = + raft::resource::comms_initialized(handle) && handle.get_comms().get_size() > 1; + this->error_checking(input, labels, n_rows, n_cols, false, true); + int device = 0; + RAFT_CUDA_TRY(cudaGetDevice(&device)); + int n_sampled_rows = 0; if (this->rf_params.bootstrap) { n_sampled_rows = std::round(this->rf_params.max_samples * n_rows); } else { @@ -135,8 +148,11 @@ class RandomForest { n_sampled_rows = n_rows; } 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()); @@ -162,6 +178,7 @@ class RandomForest { #pragma omp parallel for num_threads(n_streams) for (int i = 0; i < this->rf_params.n_trees; i++) { + RAFT_CUDA_TRY(cudaSetDevice(device)); int stream_id = omp_get_thread_num(); auto s = handle.get_stream_from_stream_pool(stream_id); diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 563553cf74..47b0acf3e5 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -207,6 +207,7 @@ if(BUILD_CUML_MG_TESTS) ConfigureTest(PREFIX MG NAME KNN_REGRESS_TEST mg/knn_regress.cu MPI RAFT_DISTRIBUTED ML_INCLUDE) ConfigureTest(PREFIX MG NAME MAIN_TEST mg/main.cu MPI RAFT_DISTRIBUTED ML_INCLUDE) ConfigureTest(PREFIX MG NAME PCA_TEST mg/pca.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..9ee6e329fa --- /dev/null +++ b/cpp/tests/mg/rf_test.cu @@ -0,0 +1,422 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * 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 + +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; + PartitionKind partition_kind; +}; + +template +void hash_value(uint64_t& hash, const T& value) +{ + const auto* bytes = reinterpret_cast(&value); + for (size_t i = 0; i < sizeof(T); ++i) { + hash ^= bytes[i]; + hash *= 1099511628211ULL; + } +} + +template +uint64_t hash_forest_structure(const RandomForestMetaData& forest) +{ + uint64_t hash = 1469598103934665603ULL; + hash_value(hash, forest.n_features); + hash_value(hash, forest.rf_params.n_trees); + hash_value(hash, forest.rf_params.bootstrap); + hash_value(hash, forest.rf_params.max_samples); + hash_value(hash, forest.rf_params.seed); + hash_value(hash, forest.rf_params.tree_params.max_depth); + hash_value(hash, forest.rf_params.tree_params.max_leaves); + hash_value(hash, forest.rf_params.tree_params.max_n_bins); + hash_value(hash, forest.rf_params.tree_params.min_samples_leaf); + hash_value(hash, forest.rf_params.tree_params.min_samples_split); + hash_value(hash, forest.rf_params.tree_params.min_impurity_decrease); + hash_value(hash, forest.rf_params.tree_params.split_criterion); + for (auto const& tree : forest.trees) { + hash_value(hash, tree->treeid); + hash_value(hash, tree->depth_counter); + hash_value(hash, tree->leaf_counter); + hash_value(hash, tree->num_outputs); + hash_value(hash, tree->sparsetree.size()); + for (auto const& node : tree->sparsetree) { + hash_value(hash, node.ColumnId()); + hash_value(hash, node.QueryValue()); + hash_value(hash, node.BestMetric()); + hash_value(hash, node.LeftChildId()); + hash_value(hash, node.InstanceCount()); + hash_value(hash, node.IsLeaf()); + } + } + return hash; +} + +template +uint64_t hash_forest_leaf_values(const RandomForestMetaData& forest) +{ + uint64_t hash = 1469598103934665603ULL; + for (auto const& tree : forest.trees) { + hash_value(hash, tree->vector_leaf.size()); + for (auto const& leaf : tree->vector_leaf) { + hash_value(hash, leaf); + } + } + return hash; +} + +template +uint64_t hash_host_vector(std::vector const& values) +{ + uint64_t hash = 1469598103934665603ULL; + hash_value(hash, values.size()); + for (auto const& value : values) { + hash_value(hash, value); + } + return hash; +} + +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) +{ + X.resize(rows.size() * params.n_cols); + y.resize(rows.size()); + 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[i * params.n_cols + col] = 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); + } + } +} + +template +void make_prediction_dataset(RfMgTestParams const& params, std::vector& X) +{ + X.resize(params.n_rows * params.n_cols); + for (int row = 0; row < params.n_rows; ++row) { + for (int col = 0; col < params.n_cols; ++col) { + X[row * params.n_cols + col] = + static_cast(((row * 5 + col * 11 + params.seed) % 101) - 50); + } + } +} + +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()) { + 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) +{ + 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); } +} + +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 n_gpus = 0; + RAFT_CUDA_TRY(cudaGetDeviceCount(&n_gpus)); + if (n_gpus < size) { + ADD_FAILURE() << "Number of GPUs is smaller than MPI ranks: ngpus=" << n_gpus + << ", nranks=" << size; + return; + } + RAFT_CUDA_TRY(cudaSetDevice(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; + make_local_dataset(params, local_rows, h_X, h_y); + + rmm::device_uvector X(h_X.size(), handle.get_stream()); + rmm::device_uvector y(h_y.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()); + + 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 forest; + if constexpr (std::is_integral_v) { + fit(handle, + &forest, + X.data(), + static_cast(local_rows.size()), + params.n_cols, + y.data(), + params.n_labels, + rf_params); + } else { + fit(handle, + &forest, + X.data(), + static_cast(local_rows.size()), + params.n_cols, + y.data(), + rf_params); + } + + expect_global_tree_counts(forest, params.n_rows); + expect_tree_limits(forest, params); + expect_identical_across_ranks(handle, hash_forest_structure(forest), "tree structure"); + expect_identical_across_ranks(handle, hash_forest_leaf_values(forest), "leaf values"); + expect_identical_predictions_across_ranks(handle, &forest); + } + + private: + void expect_identical_across_ranks(raft::handle_t const& handle, + uint64_t local_hash, + char const* label) + { + auto const& comm = handle.get_comms(); + std::vector hashes(comm.get_size()); + MPI_Allgather(&local_hash, 1, MPI_UINT64_T, hashes.data(), 1, MPI_UINT64_T, MPI_COMM_WORLD); + for (auto hash : hashes) { + EXPECT_EQ(hash, hashes.front()) << "Mismatched distributed RF " << label; + } + } + + void expect_identical_predictions_across_ranks(raft::handle_t const& handle, + RandomForestMetaData* forest) + { + std::vector h_X; + make_prediction_dataset(params, h_X); + rmm::device_uvector X(h_X.size(), handle.get_stream()); + rmm::device_uvector X_transpose(h_X.size(), handle.get_stream()); + rmm::device_uvector predictions(params.n_rows, handle.get_stream()); + raft::update_device(X.data(), h_X.data(), h_X.size(), handle.get_stream()); + raft::linalg::transpose( + handle, X.data(), X_transpose.data(), params.n_rows, params.n_cols, handle.get_stream()); + predict(handle, forest, X_transpose.data(), params.n_rows, params.n_cols, predictions.data()); + std::vector h_predictions(params.n_rows); + raft::update_host( + h_predictions.data(), predictions.data(), h_predictions.size(), handle.get_stream()); + handle.sync_stream(); + expect_identical_across_ranks(handle, hash_host_vector(h_predictions), "predictions"); + } + + 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) {} + +std::vector inputs = { + {128, 4, 1, 1.0f, 3, -1, 16, 1, 2, 0.0f, 1, 1, GINI, 7, 2, false, PartitionKind::Contiguous}, + {128, 4, 3, 0.5f, 4, 16, 32, 1, 2, 0.0f, 4, 4, ENTROPY, 11, 2, false, PartitionKind::Strided}, + {192, 6, 1, 1.0f, 5, -1, 32, 2, 4, 0.0f, 1, 1, MSE, 13, 2, false, PartitionKind::Imbalanced}, + {96, 3, 2, 1.0f, 4, 8, 8, 1, 2, 0.0f, 1, 1, GINI, 17, 2, true, PartitionKind::Imbalanced}, + {144, 5, 2, 0.8f, 4, -1, 16, 1, 2, 0.0f, 1, 1, GINI, 31, 3, false, PartitionKind::Strided}, + {160, 5, 1, 0.8f, 4, -1, 16, 1, 2, 0.0f, 1, 1, MSE, 19, 2, true, PartitionKind::Contiguous}, + {64, + 4, + 2, + 1.0f, + 4, + -1, + 16, + 1, + 2, + 0.0f, + 3, + 3, + GINI, + 23, + 2, + false, + PartitionKind::EmptyNonRootRanks}, + {80, + 5, + 2, + 0.8f, + 4, + -1, + 16, + 1, + 2, + 0.0f, + 3, + 3, + MSE, + 29, + 2, + false, + PartitionKind::EmptyNonRootRanks}}; + +INSTANTIATE_TEST_CASE_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 3484ad9bb8..efaeada9f1 100644 --- a/cpp/tests/sg/rf_test.cu +++ b/cpp/tests/sg/rf_test.cu @@ -689,6 +689,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); +} + //------------------------------------------------------------------------------------------------------------------------------------- struct QuantileTestParameters { int n_rows; diff --git a/python/cuml/cuml/dask/ensemble/base.py b/python/cuml/cuml/dask/ensemble/base.py index d150adec06..f5f427689e 100644 --- a/python/cuml/cuml/dask/ensemble/base.py +++ b/python/cuml/cuml/dask/ensemble/base.py @@ -2,15 +2,15 @@ # SPDX-License-Identifier: Apache-2.0 # -import math import warnings from collections.abc import Iterable +import cudf import cupy as cp import dask import numpy as np -import treelite -from dask.distributed import Future +from dask.distributed import Future, get_worker +from raft_dask.common.comms import Comms, get_raft_comm_state from cuml import using_output_type from cuml.dask._compat import DASK_2025_4_0 @@ -62,43 +62,29 @@ def _create_model( 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_estimators_per_worker = [ + n_estimators for _ in range(len(self.workers)) + ] self.rfs = { worker: self.client.submit( model_func, - n_estimators=self.n_estimators_per_worker[n], - random_state=seeds[n], + n_estimators=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 + return [n_estimators for _ in range(len(self.workers))] def _fit(self, model, dataset, convert_dtype, broadcast_data): # rapids-pre-commit-hooks: disable-next-line @@ -114,8 +100,18 @@ def _fit(self, model, dataset, convert_dtype, broadcast_data): stacklevel=3, ) + if broadcast_data: + warnings.warn( + "broadcast_data is ignored for distributed RandomForest " + "training because each worker participates in one global " + "tree build over its local rows.", + UserWarning, + stacklevel=3, + ) + data = DistributedDataHandler.create(dataset, client=self.client) - self.active_workers = data.workers + fit_workers = list(self.workers) + self.active_workers = fit_workers self.datatype = data.datatype labels = self.client.persist(dataset[1]) @@ -124,81 +120,61 @@ def _fit(self, model, dataset, convert_dtype, broadcast_data): else: self.num_classes = len(dask.array.unique(labels).compute()) - combined_data = ( - list(map(lambda x: x[1], data.gpu_futures)) - if broadcast_data - else None - ) + global_n_rows = sum(total for _, total in data._worker_sizes.values()) + if global_n_rows <= 0: + raise ValueError("RandomForest requires at least one global row") - futures = list() - for idx, (worker, worker_data) in enumerate( - data.worker_to_parts.items() - ): - futures.append( - self.client.submit( - _func_fit, + n_cols = dataset[0].shape[1] + x_dtype = _dtype_from_input(dataset[0]) + y_dtype = _dtype_from_input(dataset[1]) + classes = getattr(self, "unique_classes", None) + + comms = Comms( + comms_p2p=False, client=self.client, streams_per_handle=1 + ) + comms.init(workers=fit_workers) + futures = [] + fit_futures = {} + try: + for worker in fit_workers: + worker_data = data.worker_to_parts.get(worker) + fit_future = self.client.submit( + _func_fit_distributed, model[worker], - combined_data if broadcast_data else worker_data, + comms.sessionId, + worker_data, convert_dtype, + self.datatype, + x_dtype, + y_dtype, + n_cols, + global_n_rows, + classes, workers=[worker], pure=False, ) - ) - - 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) - - 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) + fit_futures[worker] = fit_future + futures.append(fit_future) + wait_and_raise_from_futures(futures) + finally: + comms.destroy() + + self.rfs.update(fit_futures) + self.n_active_estimators_per_worker = [ + self.n_estimators for _ in self.active_workers + ] + self._set_internal_model(futures[0].result()) 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. + Return one worker model. + + Distributed training now synchronizes the core tree builder across + workers, so each worker owns an equivalent full forest. The old Dask + implementation concatenated independent per-worker sub-forests here. """ - 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 + return self.rfs[self.active_workers[0]].result() def _partial_inference(self, X, op_type, delayed, **kwargs): data = DistributedDataHandler.create(X, client=self.client) @@ -341,6 +317,67 @@ def _func_fit(model, input_data, convert_dtype): return model.fit(X, y, convert_dtype=convert_dtype) +def _func_fit_distributed( + model, + session_id, + input_data, + convert_dtype, + datatype, + x_dtype, + y_dtype, + n_cols, + global_n_rows, + classes, +): + state = get_raft_comm_state(session_id, get_worker()) + handle = state["handle"] + if input_data is None: + X, y = _empty_local_data(datatype, x_dtype, y_dtype, n_cols) + else: + X = concatenate([item[0] for item in input_data]) + y = concatenate([item[1] for item in input_data]) + if hasattr(model, "_fit_with_handle"): + kwargs = { + "convert_dtype": convert_dtype, + "global_n_rows": global_n_rows, + } + if classes is not None: + kwargs["classes"] = classes + return model._fit_with_handle(X, y, handle, **kwargs) + return model.fit(X, y, convert_dtype=convert_dtype) + + +def _empty_local_data(datatype, x_dtype, y_dtype, n_cols): + if datatype == "cudf": + X = cudf.DataFrame( + {i: cp.empty(0, dtype=x_dtype) for i in range(n_cols)} + ) + y = cudf.Series(cp.empty(0, dtype=y_dtype)) + else: + X = cp.empty((0, n_cols), dtype=x_dtype, order="F") + y = cp.empty(0, dtype=y_dtype) + return X, y + + +def _dtype_from_input(data): + dtype = getattr(data, "dtype", None) + if dtype is not None: + return np.dtype(dtype) + + meta = getattr(data, "_meta", None) + dtype = getattr(meta, "dtype", None) + if dtype is not None: + return np.dtype(dtype) + + dtypes = getattr(meta, "dtypes", None) + if dtypes is not None: + if hasattr(dtypes, "iloc"): + return np.dtype(dtypes.iloc[0]) + return np.dtype(next(iter(dtypes))) + + raise TypeError(f"Could not determine dtype for {type(data)!r}") + + def _func_predict_partial(model, input_data, **kwargs): """ Whole dataset inference with part of the model (trees at disposal locally). diff --git a/python/cuml/cuml/ensemble/randomforest_common.pyx b/python/cuml/cuml/ensemble/randomforest_common.pyx index 99f5e9142d..6b02606fac 100644 --- a/python/cuml/cuml/ensemble/randomforest_common.pyx +++ b/python/cuml/cuml/ensemble/randomforest_common.pyx @@ -458,7 +458,7 @@ class BaseRandomForestModel(Base, InteropMixin): handle=get_handle(), ) - def _fit_forest(self, X, y): + def _fit_forest(self, X, y, handle=None, global_n_rows=None): cdef bool is_classifier = self._estimator_type == "classifier" cdef bool is_float32 = X.dtype == np.float32 @@ -466,6 +466,7 @@ class BaseRandomForestModel(Base, InteropMixin): cdef uintptr_t y_ptr = y.data.ptr cdef int n_rows = X.shape[0] cdef int n_cols = X.shape[1] + cdef int n_rows_for_params cdef level_enum verbose = self._verbose_level cdef int n_classes = self.n_classes_ if is_classifier else 0 @@ -507,21 +508,23 @@ class BaseRandomForestModel(Base, InteropMixin): 0 if self.random_state is None else check_random_seed(self.random_state) ) + n_rows_for_params = n_rows if global_n_rows is None else global_n_rows + 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 * n_rows_for_params) ) 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 * n_rows_for_params)) ) cdef int n_bins - if self.n_bins > n_rows: + if self.n_bins > n_rows_for_params: 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 = n_rows_for_params else: n_bins = self.n_bins @@ -543,7 +546,8 @@ class BaseRandomForestModel(Base, InteropMixin): ) cdef TreeliteModelHandle tl_handle - handle = get_handle(n_streams=self.n_streams) + if handle is None: + handle = get_handle(n_streams=self.n_streams) cdef handle_t* handle_ = handle.getHandle() # Store oob_score in C variable for nogil block @@ -636,7 +640,7 @@ class BaseRandomForestModel(Base, InteropMixin): TreeliteFreeModel(tl_handle), "Failed to free Treelite model:" ) - self._n_samples = y.shape[0] + self._n_samples = n_rows_for_params 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 a251b04865..72f6f1d832 100644 --- a/python/cuml/cuml/ensemble/randomforestclassifier.py +++ b/python/cuml/cuml/ensemble/randomforestclassifier.py @@ -264,6 +264,55 @@ def fit(self, X, y, *, convert_dtype=True) -> "RandomForestClassifier": self.n_classes_ = len(classes) return self._fit_forest(X, y) + def _fit_with_handle( + self, + X, + y, + handle, + *, + convert_dtype=True, + classes=None, + global_n_rows=None, + ) -> "RandomForestClassifier": + """ + Internal fit path used by Dask RF. The supplied handle may contain + distributed RAFT communicator state. + """ + if classes is None: + X, y, classes = check_inputs( + self, + X, + y, + dtype=("float32", "float64"), + convert_dtype=convert_dtype, + order="F", + y_dtype="int32", + return_classes=True, + ensure_min_samples=0, + reset=True, + ) + else: + X, y = check_inputs( + self, + X, + y, + dtype=("float32", "float64"), + convert_dtype=convert_dtype, + order="F", + y_dtype=None, + ensure_min_samples=0, + reset=True, + ) + classes = cp.asarray(classes) + y = cp.asarray(y) + y = cp.searchsorted(classes, y).astype(cp.int32) + classes = cp.asnumpy(classes) + self.classes_ = classes + self.n_classes_ = len(classes) + return self._fit_forest( + X, y, handle=handle, global_n_rows=global_n_rows + ) + @nvtx.annotate( message="predict RF-Classifier @randomforestclassifier.pyx", domain="cuml_python", diff --git a/python/cuml/cuml/ensemble/randomforestregressor.py b/python/cuml/cuml/ensemble/randomforestregressor.py index 556e6d9cc6..35a08727d0 100644 --- a/python/cuml/cuml/ensemble/randomforestregressor.py +++ b/python/cuml/cuml/ensemble/randomforestregressor.py @@ -220,6 +220,33 @@ def fit(self, X, y, *, convert_dtype=True) -> "RandomForestRegressor": ) return self._fit_forest(X, y) + def _fit_with_handle( + self, + X, + y, + handle, + *, + convert_dtype=True, + global_n_rows=None, + ) -> "RandomForestRegressor": + """ + Internal fit path used by Dask RF. The supplied handle may contain + distributed RAFT communicator state. + """ + X, y = check_inputs( + self, + X, + y, + dtype=("float32", "float64"), + convert_dtype=convert_dtype, + order="F", + ensure_min_samples=0, + reset=True, + ) + return self._fit_forest( + X, y, handle=handle, global_n_rows=global_n_rows + ) + @nvtx.annotate( message="predict RF-Regressor @randomforestclassifier.pyx", domain="cuml_python", diff --git a/python/cuml/tests/dask/test_dask_random_forest.py b/python/cuml/tests/dask/test_dask_random_forest.py index 2b74fa12b4..92768ac771 100644 --- a/python/cuml/tests/dask/test_dask_random_forest.py +++ b/python/cuml/tests/dask/test_dask_random_forest.py @@ -1,8 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 -import json - import cudf import cupy as cp import dask_cudf @@ -11,13 +9,11 @@ import pytest import treelite from dask.array import from_array -from dask.distributed import Client from sklearn.datasets import make_classification, make_regression from sklearn.ensemble import RandomForestClassifier as skrfc from sklearn.metrics import accuracy_score, mean_squared_error, r2_score from sklearn.model_selection import train_test_split -from cuml.dask._compat import DASK_2025_4_0 from cuml.dask.common import utils as dask_utils from cuml.dask.ensemble import RandomForestClassifier as cuRFC_mg from cuml.dask.ensemble import RandomForestRegressor as cuRFR_mg @@ -46,63 +42,10 @@ def _prep_training_data(c, X_train, y_train, partitions_per_worker): return X_train_df, y_train_df -@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 - c = Client(cluster) - kwargs = {"n_workers": -1} if DASK_2025_4_0() else {} - n_workers = len(c.scheduler_info(**kwargs)["workers"]) - - try: - X, y = make_classification( - n_samples=n_workers * 8000, - n_features=20, - n_clusters_per_class=1, - n_informative=10, - random_state=123, - n_classes=10, - ) - - X = X.astype(np.float32) - y = y.astype(np.int32) - - X_train, X_test, y_train, y_test = train_test_split( - X, y, test_size=n_workers * 500, random_state=123 - ) - - cu_rf_params = { - "n_estimators": n_workers * 25, - "max_depth": 16, - "n_bins": 256, - "random_state": 10, - } - - X_train_df, y_train_df = _prep_training_data( - c, X_train, y_train, partitions_per_worker - ) - - cuml_mod = cuRFC_mg(**cu_rf_params, ignore_empty_partitions=True) - 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() - acc_score_gpu = accuracy_score(cuml_preds_gpu, y_test) - - # Compare with sklearn baseline - sk_model = skrfc( - n_estimators=cu_rf_params["n_estimators"], - max_depth=cu_rf_params["max_depth"], - random_state=cu_rf_params["random_state"], - n_jobs=-1, - ) - sk_model.fit(X_train, y_train) - sk_preds = sk_model.predict(X_test) - sk_acc = accuracy_score(y_test, sk_preds) - - # Observed: mean=0.002, range=[0.002, 0.002], stderr=0.000 - assert acc_score_gpu >= (sk_acc - 0.07) - - finally: - c.close() +def _assert_num_trees(model, n_estimators): + treelite_bytes = model.internal_model._treelite_model_bytes + treelite_model = treelite.Model.deserialize_bytes(treelite_bytes) + assert treelite_model.num_tree == n_estimators @pytest.mark.parametrize("dtype", [np.float32, np.float64]) @@ -144,6 +87,7 @@ def test_rf_regression_dask_fil(partitions_per_worker, dtype, client): cuml_mod = cuRFR_mg(**cu_rf_params, ignore_empty_partitions=True) cuml_mod.fit(X_train_df, y_train_df) + _assert_num_trees(cuml_mod, cu_rf_params["n_estimators"]) cuml_mod_predict = cuml_mod.predict(X_test_df) cuml_mod_predict = cp.asnumpy(cp.array(cuml_mod_predict.compute())) @@ -185,6 +129,7 @@ def test_rf_classification_dask_array(partitions_per_worker, client): X_test_dask_array = from_array(X_test) cuml_mod = cuRFC_mg(**cu_rf_params) cuml_mod.fit(X_train_df, y_train_df) + _assert_num_trees(cuml_mod, cu_rf_params["n_estimators"]) cuml_mod_predict = cuml_mod.predict(X_test_dask_array).compute() acc_score = accuracy_score(cuml_mod_predict, y_test, normalize=True) @@ -254,160 +199,11 @@ def test_rf_classification_dask_fil_predict_proba( assert fil_mse <= sk_mse + 0.029 -@pytest.mark.parametrize("model_type", ["classification", "regression"]) -def test_rf_concatenation_dask(client, model_type): - n_workers = len(client.scheduler_info(n_workers=-1)["workers"]) - - X, y = make_classification( - n_samples=n_workers * 200, n_features=30, random_state=123, n_classes=2 - ) - - X = X.astype(np.float32) - if model_type == "classification": - y = y.astype(np.int32) - else: - y = y.astype(np.float32) - n_estimators = 40 - cu_rf_params = {"n_estimators": n_estimators, "max_depth": 16} - - X_df, y_df = _prep_training_data(client, X, y, partitions_per_worker=2) - - if model_type == "classification": - cu_rf_mg = cuRFC_mg(**cu_rf_params) - else: - 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 - - -@pytest.mark.parametrize("ignore_empty_partitions", [True, False]) -def test_single_input_regression(client, ignore_empty_partitions): - 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) - - -@pytest.mark.parametrize("max_depth", [1, 2, 3, 5, 10, 15, 20]) -@pytest.mark.parametrize("n_estimators", [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( - n_samples=n_samples_per_worker * n_workers, - n_features=20, - n_clusters_per_class=1, - n_informative=10, - random_state=123, - n_classes=2, - ) - X = X.astype(np.float32) - dask_model = cuRFC_mg( - max_features=1.0, - max_samples=1.0, - n_bins=16, - split_criterion=0, - min_samples_leaf=2, - random_state=23707, - n_streams=1, - n_estimators=n_estimators, - max_leaves=-1, - max_depth=max_depth, - ) - y = y.astype(np.int32) - - X_dask, y_dask = _prep_training_data(client, X, y, partitions_per_worker=2) - dask_model.fit(X_dask, y_dask) - model = dask_model.get_combined_model() - json_obj = json.loads(model.as_treelite().dump_as_json()) - - def check_count(node, nodes): - if "left_child" in node: - left = nodes[node["left_child"]] - right = nodes[node["right_child"]] - count = check_count(left, nodes) + check_count(right, nodes) - assert count == node["data_count"] - return node["data_count"] - - 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 - # Check that the data_count accumulates properly as you move up the tree - for node in nodes: - check_count(node, nodes) - - -def test_unlimited_max_depth_classifier(client): - n_workers = len(client.scheduler_info(n_workers=-1)["workers"]) - X, y = make_classification( - n_samples=n_workers * 200, n_features=10, random_state=42 - ) - X = X.astype(np.float32) - 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.fit(X_dask, y_dask) - preds = cp.asnumpy(cp.array(clf.predict(X_dask).compute())) - assert len(preds) == len(y) - - -def test_unlimited_max_depth_regressor(client): - n_workers = len(client.scheduler_info(n_workers=-1)["workers"]) - X, y = make_regression( - n_samples=n_workers * 200, n_features=10, random_state=42 - ) - X = X.astype(np.float32) - 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.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): 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": @@ -443,78 +239,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