From f76850091d76d9d3816197989eeacf67bcb86bb0 Mon Sep 17 00:00:00 2001 From: Rory Mitchell Date: Mon, 8 Jun 2026 05:56:51 -0700 Subject: [PATCH 1/2] Fix RF feature sampling retries --- .../batched-levelalgo/builder.cuh | 197 ++++++----- .../kernels/builder_kernels.cuh | 308 +++--------------- .../kernels/builder_kernels_impl.cuh | 17 +- .../batched-levelalgo/quantiles.cuh | 5 +- .../batched-levelalgo/random_utils.cuh | 18 + python/cuml/tests/test_random_forest.py | 22 ++ 6 files changed, 197 insertions(+), 370 deletions(-) diff --git a/cpp/src/decisiontree/batched-levelalgo/builder.cuh b/cpp/src/decisiontree/batched-levelalgo/builder.cuh index e9a7996b65..f71fe4bbf9 100644 --- a/cpp/src/decisiontree/batched-levelalgo/builder.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/builder.cuh @@ -19,9 +19,11 @@ #include +#include #include #include #include +#include namespace ML { namespace DT { @@ -182,7 +184,7 @@ struct Builder { int n_blks_for_cols = 10; /** Memory alignment value */ const size_t align_value = 512; - IdxT* colids; + IdxT* column_samples; /** rmm device workspace buffer */ rmm::device_uvector d_buff; /** pinned host buffer to store the trained nodes */ @@ -274,7 +276,8 @@ struct Builder { 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 + d_wsize += + calculateAlignedBytes(sizeof(IdxT) * max_batch * dataset.n_sampled_cols); // column_samples // all nodes in the tree h_wsize += // h_workload_info @@ -314,7 +317,7 @@ struct Builder { d_wspace += calculateAlignedBytes(sizeof(NodeWorkItem) * max_batch); workload_info = reinterpret_cast*>(d_wspace); d_wspace += calculateAlignedBytes(sizeof(WorkloadInfo) * max_blocks_dimx); - colids = reinterpret_cast(d_wspace); + column_samples = reinterpret_cast(d_wspace); d_wspace += calculateAlignedBytes(sizeof(IdxT) * max_batch * dataset.n_sampled_cols); RAFT_CUDA_TRY( @@ -377,96 +380,73 @@ struct Builder { raft::common::nvtx::range fun_scope("Builder::doSplit @builder.cuh [batched-levelalgo]"); // start fresh on the number of *new* nodes created in this batch RAFT_CUDA_TRY(cudaMemsetAsync(n_nodes, 0, sizeof(IdxT), builder_stream)); - initSplit(splits, work_items.size(), builder_stream); - - // 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); + const IdxT original_n_sampled_cols = dataset.n_sampled_cols; + ASSERT(original_n_sampled_cols > 0 && original_n_sampled_cols <= dataset.N, + "n_sampled_cols must be in [1, n_cols]"); + const std::size_t max_sampling_rounds = std::size_t( + (dataset.N + original_n_sampled_cols - 1) / original_n_sampled_cols); + struct HostSplit { + DataT quesval; + IdxT colid; + DataT best_metric_val; + int nLeft; + }; + static_assert(sizeof(HostSplit) == sizeof(SplitT)); + static_assert(alignof(HostSplit) == alignof(SplitT)); + + // The final split chosen for each original work item. Nodes that need + // additional feature samples are compacted in active_items, so successful + // splits must be copied back to their original batch position. + std::vector final_splits(work_items.size()); + // Current retry batch. It starts as the full batch and shrinks to only + // nodes whose sampled features did not produce a valid split. + std::vector active_items(work_items); + // active_items[i] maps back to the corresponding index in the original + // work_items/final_splits arrays. + std::vector active_to_original(work_items.size()); + for (std::size_t i = 0; i < active_to_original.size(); ++i) { + active_to_original[i] = i; + } - // do feature-sampling - if (dataset.n_sampled_cols != dataset.N) { - raft::common::nvtx::range fun_scope("feature-sampling"); - constexpr int block_threads = 128; - constexpr int max_samples_per_thread = 72; // register spillage if more than this limit - // decide if the problem size is suitable for the excess-sampling strategy. - // - // our required shared memory is a function of number of samples we'll need to sample (in - // parallel, with replacement) in excess to get 'k' uniques out of 'n' features. estimated - // static shared memory required by cub's block-wide collectives: - // max_samples_per_thread * block_threads * sizeof(IdxT) - // - // The maximum items to sample ( the constant `max_samples_per_thread` to be set at - // compile-time) is calibrated so that: - // 1. There is no register spills and accesses to global memory - // 2. The required static shared memory (ie, `max_samples_per_thread * block_threads * - // sizeof(IdxT)` does not exceed 46KB. - // - // number of samples we'll need to sample (in parallel, with replacement), to expect 'k' - // unique samples from 'n' is given by the following equation: log(1 - k/n)/log(1 - 1/n) ref: - // https://stats.stackexchange.com/questions/296005/the-expected-number-of-unique-elements-drawn-with-replacement - IdxT n_parallel_samples = - std::ceil(raft::log(1 - double(dataset.n_sampled_cols) / double(dataset.N)) / - (raft::log(1 - 1.f / double(dataset.N)))); - // maximum sampling work possible by all threads in a block : - // `max_samples_per_thread * block_thread` - // dynamically calculated sampling work to be done per block: - // `n_parallel_samples` - // former must be greater or equal to than latter for excess-sampling-based strategy - if (max_samples_per_thread * block_threads >= n_parallel_samples) { - raft::common::nvtx::range fun_scope("excess-sampling-based approach"); - dim3 grid; - grid.x = work_items.size(); - grid.y = 1; - grid.z = 1; - - if (n_parallel_samples <= block_threads) - // each thread randomly samples only 1 sample - excess_sample_with_replacement_kernel - <<>>(colids, - d_work_items, - work_items.size(), - treeid, - seed, - dataset.N, - dataset.n_sampled_cols, - n_parallel_samples); - else - // each thread does more work and samples `max_samples_per_thread` samples - excess_sample_with_replacement_kernel - <<>>(colids, - d_work_items, - work_items.size(), - treeid, - seed, - dataset.N, - dataset.n_sampled_cols, - n_parallel_samples); - raft::common::nvtx::pop_range(); - } else { - raft::common::nvtx::range fun_scope("reservoir-sampling-based approach"); - // using algo-L (reservoir sampling) strategy to sample 'dataset.n_sampled_cols' unique - // features from 'dataset.N' total features - dim3 grid; - grid.x = (work_items.size() + 127) / 128; - grid.y = 1; - grid.z = 1; - algo_L_sample_kernel<<>>( - colids, d_work_items, work_items.size(), treeid, seed, dataset.N, dataset.n_sampled_cols); - raft::common::nvtx::pop_range(); + // Match sklearn's behavior of searching beyond max_features when the + // sampled features do not yield a valid split. + for (std::size_t round = 0; !active_items.empty() && round < max_sampling_rounds; ++round) { + IdxT sample_offset = IdxT(round) * original_n_sampled_cols; + dataset.n_sampled_cols = std::min(original_n_sampled_cols, dataset.N - sample_offset); + computeBestSplits(active_items, seed, sample_offset); + + std::vector retry_items; + std::vector retry_to_original; + for (std::size_t i = 0; i < active_items.size(); ++i) { + const auto original_idx = active_to_original[i]; + final_splits[original_idx] = HostSplit{h_splits[i].quesval, + h_splits[i].colid, + h_splits[i].best_metric_val, + h_splits[i].nLeft}; + if (SplitNotValid(h_splits[i], + params.min_impurity_decrease, + params.min_samples_leaf, + active_items[i].instances.count)) { + retry_items.push_back(active_items[i]); + retry_to_original.push_back(original_idx); + } } - RAFT_CUDA_TRY(cudaPeekAtLastError()); - raft::common::nvtx::pop_range(); - } - // 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); - RAFT_CUDA_TRY(cudaPeekAtLastError()); + if (round + 1 >= max_sampling_rounds) { break; } + active_items = std::move(retry_items); + active_to_original = std::move(retry_to_original); } - - // create child nodes (or make the current ones leaf) + dataset.n_sampled_cols = original_n_sampled_cols; + + // Partition samples once, using the valid split found for each node. Nodes + // still without a valid split after all features have been visited remain leaves. + RAFT_CUDA_TRY(cudaMemcpyAsync(splits, + final_splits.data(), + sizeof(SplitT) * work_items.size(), + cudaMemcpyHostToDevice, + builder_stream)); + raft::update_device(d_work_items, work_items.data(), work_items.size(), builder_stream); raft::common::nvtx::push_range("nodeSplitKernel @builder.cuh [batched-levelalgo]"); launchNodeSplitKernel(params.min_samples_leaf, params.min_samples_split, @@ -484,6 +464,45 @@ struct Builder { return std::make_tuple(h_splits, work_items.size()); } + void computeBestSplits(const std::vector& work_items, + uint64_t sampling_seed, + IdxT sample_offset) + { + initSplit(splits, work_items.size(), builder_stream); + RAFT_CUDA_TRY(cudaMemsetAsync( + done_count, 0, sizeof(int) * params.max_batch_size * n_blks_for_cols, builder_stream)); + RAFT_CUDA_TRY(cudaMemsetAsync(mutex, 0, sizeof(int) * params.max_batch_size, builder_stream)); + 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); + + sampleFeatures(work_items, sampling_seed, sample_offset); + + for (IdxT c = 0; c < dataset.n_sampled_cols; c += n_blks_for_cols) { + computeSplit(c, n_blocks_dimx, n_large_nodes); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + raft::update_host(h_splits, splits, work_items.size(), builder_stream); + handle.sync_stream(builder_stream); + } + + void sampleFeatures(const std::vector& work_items, + uint64_t sampling_seed, + IdxT sample_offset) + { + raft::common::nvtx::range fun_scope("feature-sampling"); + sample_features( + column_samples, + d_work_items, + work_items.size(), + treeid, + sampling_seed, + sample_offset, + dataset.N, + dataset.n_sampled_cols, + builder_stream); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + auto computeSplitSmemSize() { size_t smem_size_1 = @@ -530,7 +549,7 @@ struct Builder { quantiles, d_work_items, col, - colids, + column_samples, done_count, mutex, splits, diff --git a/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh b/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh index 6f56228d35..03d9616383 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh @@ -12,10 +12,11 @@ #include -#include -#include - -#include +#include +#include +#include +#include +#include namespace ML { namespace DT { @@ -64,6 +65,42 @@ DI OutT* alignPointer(InT dataset) return reinterpret_cast(raft::alignTo(reinterpret_cast(dataset), sizeof(OutT))); } +template +void sample_features(IdxT* column_samples, + const NodeWorkItem* work_items, + size_t work_items_size, + IdxT treeid, + uint64_t seed, + IdxT sample_offset, + IdxT n, + IdxT k, + cudaStream_t stream) +{ + auto n_column_samples = work_items_size * size_t(k); + auto counting = thrust::make_counting_iterator(0); + + thrust::for_each( + thrust::cuda::par.on(stream), + counting, + counting + n_column_samples, + [=] __device__(size_t sample_idx) { + auto node_idx = sample_idx / size_t(k); + IdxT column_index = static_cast(sample_idx % size_t(k)); + + if (k == n) { + column_samples[sample_idx] = column_index; + return; + } + + const uint32_t nodeid = work_items[node_idx].idx; + uint32_t rng_seed = fnv1a32_hash(seed, treeid, nodeid); + + cuda::shuffle_iterator shuffled_features( + n, cuda::std::minstd_rand(rng_seed), sample_offset); + column_samples[sample_idx] = shuffled_features[column_index]; + }); +} + template void launchNodeSplitKernel(const IdxT min_samples_leaf, const IdxT min_samples_split, @@ -104,267 +141,6 @@ HDI IdxT lower_bound(DataT* array, IdxT len, DataT element) return start; } -template -struct CustomDifference { - IdxT n; // Threshold for valid values - __device__ CustomDifference(IdxT n_) : n(n_) {} - __device__ IdxT operator()(const IdxT& lhs, const IdxT& rhs) - { - // Filter out placeholders (value == n) - if (lhs == n || rhs == n) return 0; // Filter out placeholders - - // Both are valid features [0, n-1] or one is the initial value (UINT_MAX) - if (lhs == rhs) - return 0; - else - return 1; - } -}; - -/** - * @brief Generates 'k' unique samples of features from 'n' feature sample-space. - * Does this for each work-item (node), feeding a unique seed for each (treeid, nodeid - * (=blockIdx.x), threadIdx.x). Method used is a random, parallel, sampling with replacement of - * excess of 'k' samples (hence the name) and then eliminating the duplicates by ordering them. The - * excess number of samples (=`n_parallel_samples`) is calculated such that after ordering there is - * at least 'k' uniques. - */ -template -CUML_KERNEL void excess_sample_with_replacement_kernel( - IdxT* colids, - const NodeWorkItem* work_items, - size_t work_items_size, - IdxT treeid, - uint64_t seed, - size_t n /* total cols to sample from*/, - size_t k /* number of unique cols to sample */, - int n_parallel_samples /* number of cols to sample with replacement */ -#ifndef NDEBUG - , - unsigned long long* feature_sample_counts = - nullptr /* optional: track feature sampling for debugging */ -#endif -) -{ - if (blockIdx.x >= work_items_size) return; - - // Specialize CUB collective types for this thread block - typedef cub::BlockRadixSort BlockRadixSortT; - typedef cub::BlockAdjacentDifference BlockAdjacentDifferenceT; - typedef cub::BlockScan BlockScanT; - - // Shared memory declarations - __shared__ union TempStorage { - typename BlockRadixSortT::TempStorage sort; - typename BlockAdjacentDifferenceT::TempStorage diff; - typename BlockScanT::TempStorage scan; - } temp_storage; - __shared__ IdxT saved_random_value; - __shared__ bool random_value_saved; - __shared__ IdxT random_offset; - - const uint32_t nodeid = work_items[blockIdx.x].idx; - - uint64_t subsequence(fnv1a32_basis); - subsequence = fnv1a32(subsequence, uint32_t(threadIdx.x)); - subsequence = fnv1a32(subsequence, uint32_t(treeid)); - subsequence = fnv1a32(subsequence, uint32_t(nodeid)); - - raft::random::PCGenerator gen(seed, subsequence, uint64_t(0)); - raft::random::UniformIntDistParams uniform_int_dist_params; - - uniform_int_dist_params.start = 0; - uniform_int_dist_params.end = n; - uniform_int_dist_params.diff = - uint64_t(uniform_int_dist_params.end - uniform_int_dist_params.start); - - IdxT n_uniques = 0; - IdxT items[MAX_SAMPLES_PER_THREAD]; - IdxT col_indices[MAX_SAMPLES_PER_THREAD]; - IdxT mask[MAX_SAMPLES_PER_THREAD]; - // populate this - for (int i = 0; i < MAX_SAMPLES_PER_THREAD; ++i) - mask[i] = 0; - if (threadIdx.x == 0) { random_value_saved = false; } - __syncthreads(); - - do { - // blocked arrangement - for (int cta_sample_idx = MAX_SAMPLES_PER_THREAD * threadIdx.x, thread_local_sample_idx = 0; - thread_local_sample_idx < MAX_SAMPLES_PER_THREAD; - ++cta_sample_idx, ++thread_local_sample_idx) { - // mask of the previous iteration, if exists, is re-used here - // so previously generated unique random numbers are used. - // newly generated random numbers may or may not duplicate the previously generated ones - // but this ensures some forward progress in order to generate at least 'k' unique random - // samples. - if (mask[thread_local_sample_idx] == 0 and cta_sample_idx < n_parallel_samples) - raft::random::custom_next( - gen, &items[thread_local_sample_idx], uniform_int_dist_params, IdxT(0), IdxT(0)); - else if (mask[thread_local_sample_idx] == - 0) // indices that exceed `n_parallel_samples` will not generate - items[thread_local_sample_idx] = n; // n is outside valid range [0, n-1] - else - continue; // this case is for samples whose mask == 1 (saving previous iteration's random - // number generated) - } - - // Save first random value before sorting (only on first iteration) - if (threadIdx.x == 0 && !random_value_saved) { - saved_random_value = items[0]; - random_value_saved = true; - } - __syncthreads(); - - // collectively sort items - BlockRadixSortT(temp_storage.sort).Sort(items); - - __syncthreads(); - - // compute the mask - // compute the adjacent differences according to the functor - // TODO: Replace deprecated 'FlagHeads' with 'SubtractLeft' when it is available - // Use -1 as the initial value since it can't match any valid column index [0, n-1] - BlockAdjacentDifferenceT(temp_storage.diff) - .SubtractLeft(items, mask, CustomDifference(n), IdxT(-1)); - - __syncthreads(); - - // do a scan on the mask to get the indices for gathering - BlockScanT(temp_storage.scan).ExclusiveSum(mask, col_indices, n_uniques); - - __syncthreads(); - - } while (n_uniques < k); - - // Use random rotation to select k features uniformly from the n_uniques sorted features - // This avoids the bias of always taking the first k sorted values - // Reuses a saved random sample as offset to avoid extra RNG consumption - - // Use the saved random value (from before sorting) to derive the offset - if (threadIdx.x == 0) { - // saved_random_value is a random value in [0, n-1], use it for offset - random_offset = saved_random_value % n_uniques; - } - __syncthreads(); - - // Write items to output with circular rotation based on random_offset - // For each unique feature at position i, its rotated position is (i + random_offset) % n_uniques - // We only output features where the rotated position is < k - IdxT col_offset = k * blockIdx.x; - for (int i = 0; i < MAX_SAMPLES_PER_THREAD; ++i) { - if (mask[i]) { // mask[i] is only set for unique, valid (non-placeholder) items - IdxT rotated_pos = (col_indices[i] + random_offset) % n_uniques; - if (rotated_pos < k) { - colids[col_offset + rotated_pos] = items[i]; -#ifndef NDEBUG - // DEBUG: Track feature sampling frequencies to expose bias - if (feature_sample_counts != nullptr) { atomicAdd(&feature_sample_counts[items[i]], 1ULL); } -#endif - } - } - } -} - -// algo L of the reservoir sampling algorithm -/** - * @brief For each work item select 'k' features without replacement from 'n' features using algo-L. - -* On exit each row of the colids array will contain k random integers from the [0..n-1] range. -* -* Each thread works on single row. The parameters work_items_size, treeid and seed are -* used to initialize a unique random seed for each work item. -* -* @param colids the generated random indices, size [work_items_size, k] row major layout -* @param work_items -* @param treeid -* @param seed -* @param n total cos to sample from -* @param k number of cols to sample - * algorithm of reservoir sampling. wiki : - * https://en.wikipedia.org/wiki/Reservoir_sampling#An_optimal_algorithm - */ -template -CUML_KERNEL void algo_L_sample_kernel(int* colids, - const NodeWorkItem* work_items, - size_t work_items_size, - IdxT treeid, - uint64_t seed, - size_t n /* total cols to sample from*/, - size_t k /* cols to sample */) -{ - int tid = threadIdx.x + blockIdx.x * blockDim.x; - if (tid >= work_items_size) return; - const uint32_t nodeid = work_items[tid].idx; - uint64_t subsequence = (uint64_t(treeid) << 32) | uint64_t(nodeid); - raft::random::PCGenerator gen(seed, subsequence, uint64_t(0)); - raft::random::UniformIntDistParams uniform_int_dist_params; - uniform_int_dist_params.start = 0; - uniform_int_dist_params.end = k; - uniform_int_dist_params.diff = - uint64_t(uniform_int_dist_params.end - uniform_int_dist_params.start); - float fp_uniform_val; - IdxT int_uniform_val; - // fp_uniform_val will have a random value between 0 and 1 - gen.next(fp_uniform_val); - double W = raft::exp(raft::log(fp_uniform_val) / k); - - size_t col(0); - // initially fill the reservoir array in increasing order of cols till k - while (1) { - colids[tid * k + col] = col; - if (col == k - 1) - break; - else - ++col; - } - // randomly sample from a geometric distribution - while (col < n) { - // fp_uniform_val will have a random value between 0 and 1 - gen.next(fp_uniform_val); - col += static_cast(raft::log(fp_uniform_val) / raft::log(1 - W)) + 1; - if (col < n) { - // int_uniform_val will now have a random value between 0...k - raft::random::custom_next(gen, &int_uniform_val, uniform_int_dist_params, IdxT(0), IdxT(0)); - colids[tid * k + int_uniform_val] = col; // the bad memory coalescing here is hidden - // fp_uniform_val will have a random value between 0 and 1 - gen.next(fp_uniform_val); - W *= raft::exp(raft::log(fp_uniform_val) / k); - } - } -} - -template -CUML_KERNEL void adaptive_sample_kernel(int* colids, - const NodeWorkItem* work_items, - size_t work_items_size, - IdxT treeid, - uint64_t seed, - int N, - int M) -{ - int tid = threadIdx.x + blockIdx.x * blockDim.x; - if (tid >= work_items_size) return; - const uint32_t nodeid = work_items[tid].idx; - - uint64_t subsequence = (uint64_t(treeid) << 32) | uint64_t(nodeid); - raft::random::PCGenerator gen(seed, subsequence, uint64_t(0)); - - int selected_count = 0; - for (int i = 0; i < N; i++) { - uint32_t toss = 0; - gen.next(toss); - uint64_t lhs = uint64_t(M - selected_count); - lhs <<= 32; - uint64_t rhs = uint64_t(toss) * (N - i); - if (lhs > rhs) { - colids[tid * M + selected_count] = i; - selected_count++; - if (selected_count == M) break; - } - } -} - template & quantiles, const NodeWorkItem* work_items, IdxT colStart, - const IdxT* colids, + const IdxT* column_samples, int* done_count, int* mutex, volatile Split* splits, 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..f525f22400 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels_impl.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels_impl.cuh @@ -207,7 +207,7 @@ static __global__ void computeSplitKernel(BinT* histograms, const Quantiles quantiles, const NodeWorkItem* work_items, IdxT colStart, - const IdxT* colids, + const IdxT* column_samples, int* done_count, int* mutex, volatile Split* splits, @@ -231,13 +231,8 @@ static __global__ void computeSplitKernel(BinT* histograms, 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; - } else { - IdxT colIndex = colStart + blockIdx.y; - col = colids[nid * dataset.n_sampled_cols + colIndex]; - } + IdxT colIndex = colStart + blockIdx.y; + IdxT col = column_samples[nid * dataset.n_sampled_cols + colIndex]; // getting the n_bins for that feature int n_bins = quantiles.n_bins_array[col]; @@ -339,7 +334,7 @@ void launchComputeSplitKernel(BinT* histograms, const Quantiles& quantiles, const NodeWorkItem* work_items, IdxT colStart, - const IdxT* colids, + const IdxT* column_samples, int* done_count, int* mutex, volatile Split* splits, @@ -360,7 +355,7 @@ void launchComputeSplitKernel(BinT* histograms, quantiles, work_items, colStart, - colids, + column_samples, done_count, mutex, splits, @@ -400,7 +395,7 @@ template void launchComputeSplitKernel<_DataT, _LabelT, _IdxT, TPB_DEFAULT, _Obj const Quantiles<_DataT, _IdxT>& quantiles, const NodeWorkItem* work_items, _IdxT colStart, - const _IdxT* colids, + const _IdxT* column_samples, int* done_count, int* mutex, volatile Split<_DataT, _IdxT>* splits, diff --git a/cpp/src/decisiontree/batched-levelalgo/quantiles.cuh b/cpp/src/decisiontree/batched-levelalgo/quantiles.cuh index 362c4004be..28290043c0 100644 --- a/cpp/src/decisiontree/batched-levelalgo/quantiles.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/quantiles.cuh @@ -37,10 +37,7 @@ static __global__ void gatherUniformSampledColumnKernel( T* out, const T* data, int sample_count, int n_rows, int col, uint64_t seed) { int tid = blockIdx.x * blockDim.x + threadIdx.x; - auto col_seed = fnv1a32_basis; - col_seed = fnv1a32(col_seed, static_cast(seed)); - col_seed = fnv1a32(col_seed, static_cast(seed >> 32)); - col_seed = fnv1a32(col_seed, static_cast(col)); + auto col_seed = fnv1a32_hash(seed, col); // Sampling is with replacement. Duplicate values from sample collisions are // removed later when quantile candidates are compacted with thrust::unique. for (int sample_idx = tid; sample_idx < sample_count; sample_idx += blockDim.x * gridDim.x) { diff --git a/cpp/src/decisiontree/batched-levelalgo/random_utils.cuh b/cpp/src/decisiontree/batched-levelalgo/random_utils.cuh index 1e5a7e3a99..f424814c14 100644 --- a/cpp/src/decisiontree/batched-levelalgo/random_utils.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/random_utils.cuh @@ -30,5 +30,23 @@ HDI uint32_t fnv1a32(uint32_t hash, uint32_t txt) return hash; } +template +HDI uint32_t fnv1a32_combine(uint32_t hash, T value) +{ + hash = fnv1a32(hash, static_cast(value)); + if constexpr (sizeof(T) > sizeof(uint32_t)) { + hash = fnv1a32(hash, static_cast(static_cast(value) >> 32)); + } + return hash; +} + +template +HDI uint32_t fnv1a32_hash(Ts... values) +{ + uint32_t hash = fnv1a32_basis; + ((hash = fnv1a32_combine(hash, values)), ...); + return hash; +} + } // namespace DT } // namespace ML diff --git a/python/cuml/tests/test_random_forest.py b/python/cuml/tests/test_random_forest.py index 0c4ee0acc7..16ecfec47b 100644 --- a/python/cuml/tests/test_random_forest.py +++ b/python/cuml/tests/test_random_forest.py @@ -1032,6 +1032,28 @@ def test_max_features(max_features, sol): assert res == sol +def test_rf_feature_sampling_retries_until_valid_split(): + n_samples = 128 + n_features = 32 + X = np.zeros((n_samples, n_features), dtype=np.float32) + y = np.zeros(n_samples, dtype=np.int32) + y[n_samples // 2 :] = 1 + X[:, 0] = y + + for random_state in range(8): + clf = curfc( + n_estimators=1, + bootstrap=False, + max_depth=None, + max_features=1, + n_bins=4, + n_streams=1, + random_state=random_state, + ) + clf.fit(X, y) + assert accuracy_score(y, clf.predict(X)) == 1.0 + + def test_rf_predict_returns_int(): X, y = make_classification() From abec448bf24255a2f17df2e595999cc2001d74dc Mon Sep 17 00:00:00 2001 From: Rory Mitchell Date: Tue, 9 Jun 2026 04:15:25 -0700 Subject: [PATCH 2/2] Address RF feature sampling review comments --- .../batched-levelalgo/builder.cuh | 37 ++++++------ .../kernels/builder_kernels.cuh | 45 +++++++-------- python/cuml/tests/test_random_forest.py | 56 ++++++++++++++++++- 3 files changed, 94 insertions(+), 44 deletions(-) diff --git a/cpp/src/decisiontree/batched-levelalgo/builder.cuh b/cpp/src/decisiontree/batched-levelalgo/builder.cuh index f71fe4bbf9..de6e760329 100644 --- a/cpp/src/decisiontree/batched-levelalgo/builder.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/builder.cuh @@ -384,8 +384,8 @@ struct Builder { const IdxT original_n_sampled_cols = dataset.n_sampled_cols; ASSERT(original_n_sampled_cols > 0 && original_n_sampled_cols <= dataset.N, "n_sampled_cols must be in [1, n_cols]"); - const std::size_t max_sampling_rounds = std::size_t( - (dataset.N + original_n_sampled_cols - 1) / original_n_sampled_cols); + const std::size_t max_sampling_rounds = + std::size_t((dataset.N + original_n_sampled_cols - 1) / original_n_sampled_cols); struct HostSplit { DataT quesval; IdxT colid; @@ -419,15 +419,11 @@ struct Builder { std::vector retry_items; std::vector retry_to_original; for (std::size_t i = 0; i < active_items.size(); ++i) { - const auto original_idx = active_to_original[i]; - final_splits[original_idx] = HostSplit{h_splits[i].quesval, - h_splits[i].colid, - h_splits[i].best_metric_val, - h_splits[i].nLeft}; - if (SplitNotValid(h_splits[i], - params.min_impurity_decrease, - params.min_samples_leaf, - active_items[i].instances.count)) { + const auto original_idx = active_to_original[i]; + final_splits[original_idx] = HostSplit{ + h_splits[i].quesval, h_splits[i].colid, h_splits[i].best_metric_val, h_splits[i].nLeft}; + if (SplitPartitionNotValid( + h_splits[i], params.min_samples_leaf, active_items[i].instances.count)) { retry_items.push_back(active_items[i]); retry_to_original.push_back(original_idx); } @@ -490,16 +486,15 @@ struct Builder { IdxT sample_offset) { raft::common::nvtx::range fun_scope("feature-sampling"); - sample_features( - column_samples, - d_work_items, - work_items.size(), - treeid, - sampling_seed, - sample_offset, - dataset.N, - dataset.n_sampled_cols, - builder_stream); + sample_features(column_samples, + d_work_items, + work_items.size(), + treeid, + sampling_seed, + sample_offset, + dataset.N, + dataset.n_sampled_cols, + builder_stream); RAFT_CUDA_TRY(cudaPeekAtLastError()); } diff --git a/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh b/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh index 03d9616383..815a4c2c0f 100644 --- a/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh +++ b/cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh @@ -48,14 +48,21 @@ struct WorkloadInfo { IdxT num_blocks; // Total number of blocks that are working on the node }; +template +HDI bool SplitPartitionNotValid(const SplitT& split, IdxT min_samples_leaf, std::size_t num_rows) +{ + return split.colid == IdxT(-1) || split.nLeft < min_samples_leaf || + (IdxT(num_rows) - split.nLeft) < min_samples_leaf; +} + template HDI bool SplitNotValid(const SplitT& split, DataT min_impurity_decrease, IdxT min_samples_leaf, std::size_t num_rows) { - return split.best_metric_val <= min_impurity_decrease || split.nLeft < min_samples_leaf || - (IdxT(num_rows) - split.nLeft) < min_samples_leaf; + return split.best_metric_val <= min_impurity_decrease || + SplitPartitionNotValid(split, min_samples_leaf, num_rows); } /* Returns 'dataset' rounded up to a correctly-aligned pointer of type OutT* */ @@ -79,26 +86,20 @@ void sample_features(IdxT* column_samples, auto n_column_samples = work_items_size * size_t(k); auto counting = thrust::make_counting_iterator(0); - thrust::for_each( - thrust::cuda::par.on(stream), - counting, - counting + n_column_samples, - [=] __device__(size_t sample_idx) { - auto node_idx = sample_idx / size_t(k); - IdxT column_index = static_cast(sample_idx % size_t(k)); - - if (k == n) { - column_samples[sample_idx] = column_index; - return; - } - - const uint32_t nodeid = work_items[node_idx].idx; - uint32_t rng_seed = fnv1a32_hash(seed, treeid, nodeid); - - cuda::shuffle_iterator shuffled_features( - n, cuda::std::minstd_rand(rng_seed), sample_offset); - column_samples[sample_idx] = shuffled_features[column_index]; - }); + thrust::for_each(thrust::cuda::par.on(stream), + counting, + counting + n_column_samples, + [=] __device__(size_t sample_idx) { + auto node_idx = sample_idx / size_t(k); + IdxT column_index = static_cast(sample_idx % size_t(k)); + + const uint32_t nodeid = work_items[node_idx].idx; + uint32_t rng_seed = fnv1a32_hash(seed, treeid, nodeid); + + cuda::shuffle_iterator shuffled_features( + n, cuda::std::minstd_rand(rng_seed), sample_offset); + column_samples[sample_idx] = shuffled_features[column_index]; + }); } template diff --git a/python/cuml/tests/test_random_forest.py b/python/cuml/tests/test_random_forest.py index 16ecfec47b..6f9a7aafa4 100644 --- a/python/cuml/tests/test_random_forest.py +++ b/python/cuml/tests/test_random_forest.py @@ -1051,7 +1051,61 @@ def test_rf_feature_sampling_retries_until_valid_split(): random_state=random_state, ) clf.fit(X, y) - assert accuracy_score(y, clf.predict(X)) == 1.0 + cuml_acc = accuracy_score(y, clf.predict(X)) + + sk_clf = skrfc( + n_estimators=1, + bootstrap=False, + max_depth=None, + max_features=1, + random_state=random_state, + ) + sk_clf.fit(X, y) + sk_acc = accuracy_score(y, sk_clf.predict(X)) + + assert sk_acc == 1.0 + assert cuml_acc == sk_acc + + +def test_rf_feature_sampling_does_not_retry_below_impurity_threshold(): + n_samples = 128 + n_features = 32 + X = np.zeros((n_samples, n_features), dtype=np.float32) + y = np.zeros(n_samples, dtype=np.int32) + y[n_samples // 2 :] = 1 + + X[:, :-1] = (np.arange(n_samples) % 2).reshape(-1, 1) + X[:, -1] = y + + cuml_accs = [] + sk_accs = [] + for random_state in range(16): + clf = curfc( + n_estimators=1, + bootstrap=False, + max_depth=None, + max_features=1, + min_impurity_decrease=0.1, + n_bins=4, + n_streams=1, + random_state=random_state, + ) + clf.fit(X, y) + cuml_accs.append(accuracy_score(y, clf.predict(X))) + + sk_clf = skrfc( + n_estimators=1, + bootstrap=False, + max_depth=None, + max_features=1, + min_impurity_decrease=0.1, + random_state=random_state, + ) + sk_clf.fit(X, y) + sk_accs.append(accuracy_score(y, sk_clf.predict(X))) + + assert min(sk_accs) == 0.5 + assert min(cuml_accs) == 0.5 def test_rf_predict_returns_int():