diff --git a/cpp/src/randomforest/randomforest.cu b/cpp/src/randomforest/randomforest.cu index 551a7b078b..35f488958b 100644 --- a/cpp/src/randomforest/randomforest.cu +++ b/cpp/src/randomforest/randomforest.cu @@ -955,4 +955,11 @@ template CUML_EXPORT void fit_treelite(const raft::handle_t& use rapids_logger::level_enum verbosity, const double* sample_weight, bool input_row_major); + +// Explicit instantiations for the implementation declared in randomforest.cuh. +template class RandomForest; +template class RandomForest; +template class RandomForest; +template class RandomForest; + } // End namespace ML diff --git a/cpp/src/randomforest/randomforest.cuh b/cpp/src/randomforest/randomforest.cuh index a19916a4a3..1365737b69 100644 --- a/cpp/src/randomforest/randomforest.cuh +++ b/cpp/src/randomforest/randomforest.cuh @@ -43,8 +43,12 @@ #define omp_get_max_threads() 1 #endif +#include #include #include +#include +#include +#include namespace ML { namespace detail { @@ -69,18 +73,22 @@ class RowSampler { int n_streams, bool* bootstrap_masks, const double* sample_weight) - : bootstrap_(rf_params.bootstrap), + : handle_(handle), + bootstrap_(rf_params.bootstrap), seed_(rf_params.seed), n_rows_(n_rows), n_sampled_rows_(n_sampled_rows), bootstrap_masks_(bootstrap_masks), sample_weight_(sample_weight), + distributed_(raft::resource::comms_initialized(handle) && handle.get_comms().get_size() > 1), + rank_(distributed_ ? handle.get_comms().get_rank() : 0), + comm_size_(distributed_ ? handle.get_comms().get_size() : 1), sample_weight_sum_(0.0), sample_weight_cdf_(0, handle.get_stream()) { ASSERT(bootstrap_masks_ == nullptr || DT::is_dev_ptr(bootstrap_masks_), "bootstrap_masks must be a GPU pointer"); - validate_sample_weight(handle, sample_weight_, n_rows_); + validate_sample_weight(handle); if (use_weighted_bootstrap()) { sample_weight_cdf_.resize(ML::narrow_cast(n_rows_), handle.get_stream()); thrust::inclusive_scan(rmm::exec_policy(handle.get_stream()), @@ -89,21 +97,17 @@ class RowSampler { sample_weight_cdf_.begin()); } - // Empty distributed partitions still pass a non-null pointer so every rank selects the same - // weighted objective type, but they have no local weight sum to validate. - if (sample_weight_ != nullptr && n_rows_ > 0) { - sample_weight_sum_ = compute_sample_weight_sum(handle); + if (sample_weight_ != nullptr) { + compute_global_sample_weights(handle); ASSERT(sample_weight_sum_ > 0.0, "sample_weight values must contain at least one positive value"); } // Use a deque instead of vector because device_uvector has a deleted copy constructor. - auto const n_sampled_rows_size = ML::narrow_cast(n_sampled_rows_); + // Place empty buffers for now; we will resize them later. for (int i = 0; i < n_streams; i++) { auto stream = handle.get_stream_from_stream_pool(i); - selected_rows_.emplace_back(n_sampled_rows_size, stream); - if (use_weighted_bootstrap()) { - weighted_draw_scratch_.emplace_back(n_sampled_rows_size, stream); - } + selected_rows_.emplace_back(0, stream); + if (use_weighted_bootstrap()) { weighted_draw_scratch_.emplace_back(0, stream); } } } @@ -114,16 +118,36 @@ class RowSampler { { raft::common::nvtx::range fun_scope("bootstrapping row IDs @randomforest.cuh"); - auto& selected_rows = selected_rows_[stream_id]; - if (n_rows_ == 0) { return selected_rows; } + auto& selected_rows = selected_rows_[stream_id]; + std::int64_t tree_sample_count = n_sampled_rows_; + if (use_weighted_bootstrap()) { + tree_sample_count = compute_rank_sample_counts(tree_id, stream); + auto tree_sample_count_size = ML::narrow_cast(tree_sample_count); + selected_rows.resize(tree_sample_count_size, stream); + weighted_draw_scratch_[stream_id].resize(tree_sample_count_size, stream); + } else { + auto tree_sample_count_size = ML::narrow_cast(tree_sample_count); + selected_rows.resize(tree_sample_count_size, stream); + } + + if (n_rows_ == 0 || tree_sample_count == 0) { + selected_rows.resize(0, stream); + store_bootstrap_mask(tree_id, selected_rows, stream); + return selected_rows; + } raft::resources stream_resources; raft::resource::set_cuda_stream(stream_resources, stream); // Hash these together so per-tree row samples are uncorrelated. - auto rs = DT::fnv1a32_basis; - rs = DT::fnv1a32(rs, seed_); - rs = DT::fnv1a32(rs, tree_id); + std::uint32_t rs; + if (distributed_ && use_weighted_bootstrap()) { + // Rank-local draws must use independent random streams. The multinomial allocation already + // determines how many draws each rank owns. + rs = DT::fnv1a32_hash(seed_, tree_id, rank_, kLocalSamplingRngTag); + } else { + rs = DT::fnv1a32_hash(seed_, tree_id); + } raft::random::RngState rng_state(rs, raft::random::GenPhilox); if (use_weighted_bootstrap()) { @@ -134,7 +158,7 @@ class RowSampler { weighted_draw_scratch.data(), weighted_draw_scratch.size(), 0.0, - sample_weight_sum_); + rank_sample_weight_sums_[rank_]); thrust::upper_bound(rmm::exec_policy(stream), sample_weight_cdf_.data(), sample_weight_cdf_.data() + n_rows_, @@ -156,7 +180,8 @@ class RowSampler { selected_rows.begin(), NonzeroSampleWeight{}); auto n_selected = selected_rows_end - selected_rows.begin(); - ASSERT(n_selected > 0, "sample_weight values must contain at least one positive value"); + ASSERT(distributed_ || n_selected > 0, + "sample_weight values must contain at least one positive value"); selected_rows.resize(n_selected, stream); } else { selected_rows.resize(ML::narrow_cast(n_sampled_rows_), stream); @@ -171,6 +196,9 @@ class RowSampler { const double* tree_sample_weight() const { return bootstrap_ ? nullptr : sample_weight_; } private: + static constexpr std::uint32_t kLocalSamplingRngTag = 0x4c4f4341; + static constexpr std::uint32_t kRankAllocationRngTag = 0x414c4c4f; + void store_bootstrap_mask(int tree_id, rmm::device_uvector& selected_rows, cudaStream_t stream) @@ -186,44 +214,176 @@ class RowSampler { tree_mask); } - double compute_sample_weight_sum(const raft::handle_t& handle) const + // Determine sample counts for all ranks by drawing from the multinomial distribution + // Multinomial(M; p_1, ..., p_P), where the parameters are: + // M = n_sampled_rows_ (total sample count), + // p_i = S_i / S, + // S_i = sum of sample weights in rank i, + // S = sum of sample weights in all ranks. + // Rank 0 will draw from the multinomial distribution (comm_size_) times, and then + // broadcast the sample count to all ranks. The function will then return the rank-local + // sample count. + std::int64_t compute_rank_sample_counts(int tree_id, cudaStream_t stream) { - if (use_weighted_bootstrap()) { - double weight_sum = 0.0; - raft::update_host( - &weight_sum, sample_weight_cdf_.data() + n_rows_ - 1, 1, handle.get_stream()); - handle.sync_stream(); - return weight_sum; + if (!distributed_) { return n_sampled_rows_; } + + rmm::device_uvector device_counts(ML::narrow_cast(comm_size_), + stream); + if (rank_ == 0) { + std::vector host_counts(comm_size_, 0); + auto rs = DT::fnv1a32_hash(seed_, tree_id, kRankAllocationRngTag); + std::mt19937_64 generator(rs); + + int last_positive_rank = comm_size_ - 1; + while (rank_sample_weight_sums_[last_positive_rank] == 0.0) { + --last_positive_rank; + } + + auto remaining_count = n_sampled_rows_; + auto remaining_weight = sample_weight_sum_; + for (int current_rank = 0; current_rank < last_positive_rank; ++current_rank) { + auto rank_weight = rank_sample_weight_sums_[current_rank]; + auto following_weight = std::max(0.0, remaining_weight - rank_weight); + if (remaining_count > 0 && rank_weight > 0.0) { + if (following_weight == 0.0) { + host_counts[current_rank] = remaining_count; + } else { + auto probability = std::clamp(rank_weight / remaining_weight, 0.0, 1.0); + std::binomial_distribution distribution(remaining_count, probability); + host_counts[current_rank] = distribution(generator); + } + } + remaining_count -= host_counts[current_rank]; + remaining_weight = following_weight; + } + host_counts[last_positive_rank] = remaining_count; + raft::update_device(device_counts.data(), host_counts.data(), host_counts.size(), stream); } - return thrust::reduce( - rmm::exec_policy(handle.get_stream()), sample_weight_, sample_weight_ + n_rows_, 0.0); + handle_.get_comms().bcast(device_counts.data(), device_counts.size(), 0, stream); + ASSERT(handle_.get_comms().sync_stream(stream) == raft::comms::status_t::SUCCESS, + "An error occurred in the distributed RF sample-count broadcast."); + + std::int64_t local_count = 0; + raft::update_host(&local_count, device_counts.data() + rank_, 1, stream); + handle_.sync_stream(stream); + return local_count; } - static void validate_sample_weight(const raft::handle_t& handle, - const double* sample_weight, - std::int64_t n_rows) + // Compute the sum of sample weights per rank and gather these per-rank sums into + // rank_sample_weight_sums_. Also compute sample_weight_sum_, the sum of sample weights over all + // ranks. + void compute_global_sample_weights(const raft::handle_t& handle) { - ASSERT(sample_weight == nullptr || DT::is_dev_ptr(sample_weight), - "sample_weight must be a GPU pointer"); - if (sample_weight == nullptr) { return; } - - bool has_invalid = thrust::any_of(rmm::exec_policy(handle.get_stream()), - sample_weight, - sample_weight + n_rows, - InvalidSampleWeight{}); - ASSERT(!has_invalid, "sample_weight values must be finite and non-negative"); + double local_sample_weight_sum{0.0}; + auto stream = handle.get_stream().get(); + if (n_rows_ > 0) { + if (use_weighted_bootstrap()) { + raft::update_host( + &local_sample_weight_sum, sample_weight_cdf_.data() + n_rows_ - 1, 1, stream); + handle.sync_stream(stream); + } else { + local_sample_weight_sum = + thrust::reduce(rmm::exec_policy(stream), sample_weight_, sample_weight_ + n_rows_, 0.0); + } + } + + if (!distributed_) { + sample_weight_sum_ = local_sample_weight_sum; + rank_sample_weight_sums_ = {local_sample_weight_sum}; + return; + } + + rmm::device_uvector local_weight_sum(1, stream); + rmm::device_uvector rank_weight_sums(comm_size_, stream); + raft::update_device(local_weight_sum.data(), &local_sample_weight_sum, 1, stream); + handle.get_comms().allgather(local_weight_sum.data(), rank_weight_sums.data(), 1, stream); + ASSERT(handle.get_comms().sync_stream(stream) == raft::comms::status_t::SUCCESS, + "An error occurred in the distributed RF weight-sum allgather."); + + rank_sample_weight_sums_.resize(comm_size_); + raft::update_host( + rank_sample_weight_sums_.data(), rank_weight_sums.data(), rank_weight_sums.size(), stream); + handle.sync_stream(stream); + sample_weight_sum_ = + std::accumulate(rank_sample_weight_sums_.begin(), rank_sample_weight_sums_.end(), 0.0); + } + + void validate_sample_weight(const raft::handle_t& handle) const + { + ASSERT(n_rows_ >= 0, "n_rows must be non-negative"); + ASSERT(n_sampled_rows_ >= 0, "n_sampled_rows must be non-negative"); + + ASSERT(sample_weight_ == nullptr || DT::is_dev_ptr(sample_weight_), + "sample_weight_ must be a GPU pointer"); + + if (distributed_) { + // All ranks must have consistent sample_weight_. + // Either all ranks have non-null sample_weight_, or + // all ranks must have sample_weight_ set to null. + // Also, ensure that n_sampled_rows_ is identical for all ranks. + auto stream = handle.get_stream().get(); + rmm::device_uvector local_values(2, stream); + rmm::device_uvector gathered_values(ML::checked_mul(2, comm_size_), + stream); + std::int64_t h_local_values[2] = {sample_weight_ == nullptr ? 0 : 1, n_sampled_rows_}; + raft::update_device(local_values.data(), h_local_values, 2, stream); + handle.get_comms().allgather(local_values.data(), gathered_values.data(), 2, stream); + ASSERT(handle.get_comms().sync_stream(stream) == raft::comms::status_t::SUCCESS, + "An error occurred while validating distributed RF sample weights."); + + std::vector h_gathered_values(gathered_values.size()); + raft::update_host( + h_gathered_values.data(), gathered_values.data(), gathered_values.size(), stream); + handle.sync_stream(stream); + for (int current_rank = 0; current_rank < comm_size_; ++current_rank) { + ASSERT(h_gathered_values[2 * current_rank] == h_gathered_values[0], + "sample_weight must be supplied consistently on every rank"); + if (bootstrap_ && sample_weight_ != nullptr) { + ASSERT(h_gathered_values[2 * current_rank + 1] == h_gathered_values[1], + "n_sampled_rows must be identical on every rank for weighted bootstrapping"); + } + } + } + + // Ensure that every element in sample_weight is non-negative and finite. + bool has_invalid_sample_weight = false; + if (sample_weight_) { + has_invalid_sample_weight = thrust::any_of(rmm::exec_policy(handle.get_stream()), + sample_weight_, + sample_weight_ + n_rows_, + InvalidSampleWeight{}); + } + if (distributed_) { + auto stream = handle.get_stream().get(); + rmm::device_uvector sample_invalid_flag(1, stream); + sample_invalid_flag.set_element(0, has_invalid_sample_weight ? 1 : 0, stream); + handle.get_comms().allreduce( + sample_invalid_flag.data(), sample_invalid_flag.data(), 1, raft::comms::op_t::MAX, stream); + ASSERT(handle.get_comms().sync_stream(stream) == raft::comms::status_t::SUCCESS, + "An error occurred while validating distributed RF sample weights."); + std::int32_t h_sample_invalid_flag; + raft::update_host(&h_sample_invalid_flag, sample_invalid_flag.data(), 1, stream); + handle.sync_stream(stream); + has_invalid_sample_weight = h_sample_invalid_flag != 0; + } + ASSERT(!has_invalid_sample_weight, "sample_weight values must be finite and non-negative"); } bool use_weighted_bootstrap() const { return bootstrap_ && sample_weight_ != nullptr; } + const raft::handle_t& handle_; bool bootstrap_; uint64_t seed_; std::int64_t n_rows_; std::int64_t n_sampled_rows_; bool* bootstrap_masks_; const double* sample_weight_; + bool distributed_; + int rank_; + int comm_size_; double sample_weight_sum_; + std::vector rank_sample_weight_sums_; rmm::device_uvector sample_weight_cdf_; std::deque> selected_rows_; std::deque> weighted_draw_scratch_; @@ -311,10 +471,26 @@ class RandomForest { raft::resource::comms_initialized(handle) && handle.get_comms().get_size() > 1; this->error_checking(input, labels, n_rows, n_cols, false, distributed); std::int64_t const n_rows_i64 = n_rows; - std::int64_t n_sampled_rows = 0; + std::int64_t global_n_rows = n_rows_i64; + if (distributed && this->rf_params.bootstrap) { + rmm::device_uvector device_global_n_rows(1, handle.get_stream()); + raft::update_device(device_global_n_rows.data(), &global_n_rows, 1, handle.get_stream()); + handle.get_comms().allreduce(device_global_n_rows.data(), + device_global_n_rows.data(), + 1, + raft::comms::op_t::SUM, + handle.get_stream().get()); + ASSERT( + handle.get_comms().sync_stream(handle.get_stream().get()) == raft::comms::status_t::SUCCESS, + "An error occurred in the distributed RF global row-count all-reduce."); + raft::update_host(&global_n_rows, device_global_n_rows.data(), 1, handle.get_stream()); + handle.sync_stream(); + } + std::int64_t n_sampled_rows = 0; if (this->rf_params.bootstrap) { + auto sample_population = distributed && sample_weight != nullptr ? global_n_rows : n_rows_i64; n_sampled_rows = - static_cast(std::round(this->rf_params.max_samples * n_rows_i64)); + static_cast(std::round(this->rf_params.max_samples * sample_population)); } else { if (this->rf_params.max_samples != 1.0) { CUML_LOG_WARN( @@ -507,10 +683,11 @@ class RandomForest { } }; -// class specializations -template class RandomForest; -template class RandomForest; -template class RandomForest; -template class RandomForest; +// Instantiated in randomforest.cu. Keep consumers of this implementation header from +// instantiating the full forest implementation and its decision-tree kernel dependencies. +extern template class RandomForest; +extern template class RandomForest; +extern template class RandomForest; +extern template class RandomForest; } // End namespace ML diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index bc6d0dbc63..12adbf397b 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -220,6 +220,9 @@ if(BUILD_CUML_MG_TESTS) ConfigureTest( PREFIX MG NAME RF_QUANTILE_TEST mg/rf_quantile_test.cu MPI RAFT_DISTRIBUTED ML_INCLUDE ) + ConfigureTest( + PREFIX MG NAME RF_ROW_SAMPLER_TEST mg/rf_row_sampler.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}'") diff --git a/cpp/tests/mg/rf_row_sampler.cu b/cpp/tests/mg/rf_row_sampler.cu new file mode 100644 index 0000000000..13ddc91f19 --- /dev/null +++ b/cpp/tests/mg/rf_row_sampler.cu @@ -0,0 +1,224 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "rf_test_utils.hpp" +#include "test_opg_utils.h" + +#include +#include +#include + +#include +#include + +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace ML { +namespace Test { +namespace opg { + +enum class WeightKind { ThreeRows, Uniform, WithZeros, Skewed }; + +struct RowSamplerTestParams { + PartitionKind partition_kind; + WeightKind weight_kind; +}; + +void initialize_mpi_once() +{ + int mpi_initialized = 0; + MPI_Initialized(&mpi_initialized); + if (!mpi_initialized) { MPI_Init(nullptr, nullptr); } +} + +void get_mpi_local_rank_size(int& local_rank, int& local_size) +{ + MPI_Comm local_comm{}; + MPI_Comm_split_type(MPI_COMM_WORLD, MPI_COMM_TYPE_SHARED, 0, MPI_INFO_NULL, &local_comm); + MPI_Comm_rank(local_comm, &local_rank); + MPI_Comm_size(local_comm, &local_size); + MPI_Comm_free(&local_comm); +} + +std::vector make_weights(WeightKind kind) +{ + switch (kind) { + case WeightKind::ThreeRows: return {0.1, 0.4, 0.5}; + case WeightKind::Uniform: return std::vector(23, 1.0); + case WeightKind::WithZeros: + return {0.0, 1.0, 0.0, 2.0, 4.0, 0.0, 3.0, 0.0, 5.0, 1.0, 0.0, 4.0, 0.0}; + case WeightKind::Skewed: + return {0.001, 0.002, 0.004, 0.008, 0.016, 0.032, 0.064, 0.128, 0.245, 0.5}; + } + return {}; +} + +class RfMgRowSamplerTest : public ::testing::TestWithParam {}; + +TEST_P(RfMgRowSamplerTest, SamplesGlobalWeightDistribution) +{ + initialize_mpi_once(); + int rank = 0; + int size = 1; + MPI_Comm_rank(MPI_COMM_WORLD, &rank); + MPI_Comm_size(MPI_COMM_WORLD, &size); + + int local_rank = 0; + int local_size = 1; + get_mpi_local_rank_size(local_rank, local_size); + int n_gpus = 0; + RAFT_CUDA_TRY(cudaGetDeviceCount(&n_gpus)); + int insufficient_local_gpus = n_gpus < local_size; + int any_insufficient_gpus = 0; + MPI_Allreduce( + &insufficient_local_gpus, &any_insufficient_gpus, 1, MPI_INT, MPI_MAX, MPI_COMM_WORLD); + if (any_insufficient_gpus) { + if (insufficient_local_gpus) { + ADD_FAILURE() << "Number of GPUs is smaller than local MPI ranks: ngpus=" << n_gpus + << ", local_ranks=" << local_size; + } + return; + } + RAFT_CUDA_TRY(cudaSetDevice(local_rank)); + + auto stream_pool = std::make_shared(1); + raft::handle_t handle(cuda::stream_ref{cudaStreamPerThread}, stream_pool); + raft::comms::initialize_mpi_comms(&handle, MPI_COMM_WORLD); + + auto global_weights = make_weights(GetParam().weight_kind); + auto local_rows = local_rows_for_rank( + static_cast(global_weights.size()), rank, size, GetParam().partition_kind); + std::vector local_weights(local_rows.size()); + std::transform(local_rows.begin(), local_rows.end(), local_weights.begin(), [&](int row) { + return global_weights[row]; + }); + + auto weight_buffer_size = std::max(std::size_t{1}, local_weights.size()); + rmm::device_uvector d_weights(weight_buffer_size, handle.get_stream()); + raft::update_device( + d_weights.data(), local_weights.data(), local_weights.size(), handle.get_stream()); + + RF_params rf_params{}; + rf_params.bootstrap = true; + rf_params.seed = 123456789ULL; + rf_params.n_streams = 1; + + constexpr std::int64_t sample_count = 200000; + detail::RowSampler sampler(handle, + rf_params, + static_cast(local_rows.size()), + sample_count, + 1, + nullptr, + d_weights.data()); + + auto stream = handle.get_stream_from_stream_pool(0); + auto& row_ids = sampler.sample(0, 0, stream.get()); + std::vector h_row_ids(row_ids.size()); + raft::update_host(h_row_ids.data(), row_ids.data(), row_ids.size(), stream); + handle.sync_stream(stream); + + std::vector local_counts(global_weights.size(), 0); + std::uint64_t local_invalid_row_count = 0; + for (auto local_row : h_row_ids) { + if (local_row < 0 || local_row >= static_cast(local_rows.size())) { + local_invalid_row_count++; + continue; + } + local_counts[local_rows[local_row]]++; + } + + std::uint64_t global_invalid_row_count = 0; + MPI_Allreduce( + &local_invalid_row_count, &global_invalid_row_count, 1, MPI_UINT64_T, MPI_SUM, MPI_COMM_WORLD); + std::vector global_counts(global_weights.size(), 0); + MPI_Allreduce(local_counts.data(), + global_counts.data(), + static_cast(global_counts.size()), + MPI_UINT64_T, + MPI_SUM, + MPI_COMM_WORLD); + + ASSERT_EQ(global_invalid_row_count, 0); + auto observed_total = + std::accumulate(global_counts.begin(), global_counts.end(), std::uint64_t{0}); + ASSERT_EQ(observed_total, static_cast(sample_count)); + + double weight_sum = std::accumulate(global_weights.begin(), global_weights.end(), 0.0); + for (std::size_t row = 0; row < global_weights.size(); ++row) { + double expected_probability = global_weights[row] / weight_sum; + double observed_probability = static_cast(global_counts[row]) / sample_count; + if (expected_probability == 0.0) { + EXPECT_EQ(global_counts[row], 0) << "global row " << row; + continue; + } + double standard_error = + std::sqrt(expected_probability * (1.0 - expected_probability) / sample_count); + double tolerance = std::max(1.0 / sample_count, 6.0 * standard_error); + EXPECT_NEAR(observed_probability, expected_probability, tolerance) << "global row " << row; + } +} + +std::string row_sampler_test_name(::testing::TestParamInfo const& test_info) +{ + char const* partition_name = nullptr; + switch (test_info.param.partition_kind) { + case PartitionKind::Contiguous: partition_name = "Contiguous"; break; + case PartitionKind::Strided: partition_name = "Strided"; break; + case PartitionKind::Imbalanced: partition_name = "Imbalanced"; break; + case PartitionKind::EmptyNonRootRanks: partition_name = "EmptyNonRootRanks"; break; + } + char const* weight_name = nullptr; + switch (test_info.param.weight_kind) { + case WeightKind::ThreeRows: weight_name = "ThreeRows"; break; + case WeightKind::Uniform: weight_name = "Uniform"; break; + case WeightKind::WithZeros: weight_name = "WithZeros"; break; + case WeightKind::Skewed: weight_name = "Skewed"; break; + } + return std::string{partition_name} + weight_name; +} + +std::vector make_row_sampler_inputs() +{ + std::vector inputs; + for (auto partition : {PartitionKind::Contiguous, + PartitionKind::Strided, + PartitionKind::Imbalanced, + PartitionKind::EmptyNonRootRanks}) { + for (auto weights : + {WeightKind::ThreeRows, WeightKind::Uniform, WeightKind::WithZeros, WeightKind::Skewed}) { + inputs.push_back({partition, weights}); + } + } + return inputs; +} + +INSTANTIATE_TEST_SUITE_P(RfRowSamplerTests, + RfMgRowSamplerTest, + ::testing::ValuesIn(make_row_sampler_inputs()), + row_sampler_test_name); + +} // 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/mg/rf_test.cu b/cpp/tests/mg/rf_test.cu index 1711294dda..f19f40ab29 100644 --- a/cpp/tests/mg/rf_test.cu +++ b/cpp/tests/mg/rf_test.cu @@ -4,6 +4,7 @@ */ #include "../prims/test_utils.h" +#include "rf_test_utils.hpp" #include "test_opg_utils.h" #include @@ -33,8 +34,6 @@ namespace ML { namespace Test { namespace opg { -enum class PartitionKind { Contiguous, Strided, Imbalanced, EmptyNonRootRanks }; - struct RfMgTestParams { int n_rows; int n_cols; @@ -134,43 +133,6 @@ void expect_forests_equal(RandomForestMetaData const& distributed_forest, } } -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, diff --git a/cpp/tests/mg/rf_test_utils.hpp b/cpp/tests/mg/rf_test_utils.hpp new file mode 100644 index 0000000000..cf351ced1f --- /dev/null +++ b/cpp/tests/mg/rf_test_utils.hpp @@ -0,0 +1,57 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +namespace ML { +namespace Test { +namespace opg { + +enum class PartitionKind { Contiguous, Strided, Imbalanced, EmptyNonRootRanks }; + +inline 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; +} + +} // namespace opg +} // namespace Test +} // namespace ML diff --git a/python/cuml/cuml/dask/ensemble/base.py b/python/cuml/cuml/dask/ensemble/base.py index 19f16a3f9a..7477e54d5e 100644 --- a/python/cuml/cuml/dask/ensemble/base.py +++ b/python/cuml/cuml/dask/ensemble/base.py @@ -149,6 +149,11 @@ def _func_fit( handle = get_raft_comm_state(session_id, get_worker())["handle"] X = concatenate([item[0] for item in input_data]) y = concatenate([item[1] for item in input_data]) + sample_weight = ( + concatenate([item[2] for item in input_data]) + if len(input_data[0]) == 3 + else None + ) model._raft_handle = handle model._distributed_n_rows = total_rows if classes is not None: @@ -156,22 +161,7 @@ def _func_fit( if class_counts is not None: model._distributed_class_counts = class_counts try: - validation_error = None - try: - _, _, sample_weight = model._prepare_fit_inputs(X, y) - if sample_weight is not None and sample_weight.sum().item() <= 0.0: - raise ValueError( - "Rank-local sample weights must sum to a positive value" - ) - except Exception as error: - validation_error = error - - if model._allreduce_validation_status(validation_error is not None): - if validation_error is not None: - raise validation_error - raise RuntimeError("Input validation failed on another worker") - - return model.fit(X, y) + return model.fit(X, y, sample_weight=sample_weight) finally: del model._raft_handle del model._distributed_n_rows diff --git a/python/cuml/cuml/dask/ensemble/randomforestclassifier.py b/python/cuml/cuml/dask/ensemble/randomforestclassifier.py index 40ac219c9d..b9b35354bc 100755 --- a/python/cuml/cuml/dask/ensemble/randomforestclassifier.py +++ b/python/cuml/cuml/dask/ensemble/randomforestclassifier.py @@ -50,8 +50,12 @@ class RandomForestClassifier( sample with replacement. * If ``False``, the whole dataset is used to build each tree. - Weighted bootstrapping through ``sample_weight`` or ``class_weight`` - is not yet supported for distributed random forests. + With weights, rows are sampled with probabilities proportional to + their weights across the complete distributed dataset. + class_weight : dict, "balanced" or None (default = None) + Weights associated with classes. ``"balanced"`` uses inverse global + class frequencies. Class weights are multiplied by ``sample_weight`` + when both are provided. max_samples : float (default = 1.0) Ratio of dataset rows used while fitting each tree. max_depth : int or None (default = None) @@ -165,35 +169,42 @@ def fit(self, X, y, sample_weight=None): y : Dask cuDF dataframe or CuPy backed Dask Array (n_rows, 1) Labels of training examples. **y must be partitioned the same way as X** - sample_weight : array-like, optional - Sample weights are not yet supported by distributed random - forests. + sample_weight : Dask cuDF Series or CuPy backed Dask Array, optional + Non-negative weights of training examples, partitioned the same + way as X and y. With ``bootstrap=True``, weights determine + sampling probabilities. Otherwise, weights are used in the + impurity calculation. """ - if self.kwargs.get("bootstrap", True) and ( - sample_weight is not None - or self.kwargs.get("class_weight") is not None + if sample_weight is not None and ( + sample_weight.npartitions != X.npartitions + or sample_weight.npartitions != y.npartitions ): - raise NotImplementedError( - "Weighted bootstrapping is not yet supported for distributed " - "random forests. Set bootstrap=False or do not provide " - "sample_weight or class_weight." - ) - if sample_weight is not None: - raise NotImplementedError( - "sample_weight is not yet supported for distributed random " - "forests" + raise ValueError( + "X, y, and sample_weight must have the same number of " + "partitions" ) + if isinstance(y, dask.array.Array): + y_for_class_counts = y + if len(y.shape) == 2: + if y.shape[1] != 1: + raise ValueError( + "y must have exactly one column when two-dimensional" + ) + y_for_class_counts = y.flatten() + # Dask implements ``unique(return_counts=True)`` using structured # arrays, which CuPy does not support. Compute the unique labels # first, then count all labels together with scalar reductions. - unique_vals = cp.asarray(dask.array.unique(y).compute()) + unique_vals = cp.asarray( + dask.array.unique(y_for_class_counts).compute() + ) if unique_vals.size == 0: class_counts = cp.empty(0, dtype=cp.int64) else: class_counts = dask.array.stack( [ - (y == class_value).sum() + (y_for_class_counts == class_value).sum() for class_value in cp.asnumpy(unique_vals) ] ).compute() @@ -202,14 +213,35 @@ def fit(self, X, y, sample_weight=None): classes = cp.asnumpy(unique_vals[order]) class_counts = cp.asnumpy(class_counts[order]) else: + y_for_class_counts = y counts_by_class = y.value_counts().compute().sort_index() classes = cp.asnumpy(cp.asarray(counts_by_class.index)) class_counts = cp.asnumpy(cp.asarray(counts_by_class)) + + if ( + sample_weight is not None + and self.kwargs.get("class_weight") == "balanced" + ): + # Use weighted class frequencies if class_weight="balanced" + # and sample_weight is provided. + class_counts = cp.asnumpy( + cp.asarray( + dask.compute( + *[ + ( + sample_weight + * (y_for_class_counts == class_value) + ).sum() + for class_value in classes + ] + ) + ) + ) self.classes_ = classes self._set_internal_model(None) self._fit( model=self.rfs, - dataset=(X, y), + dataset=(X, y) if sample_weight is None else (X, y, sample_weight), classes=classes, class_counts=class_counts, ) diff --git a/python/cuml/cuml/dask/ensemble/randomforestregressor.py b/python/cuml/cuml/dask/ensemble/randomforestregressor.py index 9a43da5727..2347fec4b4 100755 --- a/python/cuml/cuml/dask/ensemble/randomforestregressor.py +++ b/python/cuml/cuml/dask/ensemble/randomforestregressor.py @@ -39,8 +39,8 @@ class RandomForestRegressor( sample with replacement. * If ``False``, the whole dataset is used to build each tree. - Weighted bootstrapping through ``sample_weight`` is not yet supported - for distributed random forests. + With weights, rows are sampled with probabilities proportional to + their weights across the complete distributed dataset. max_samples : float (default = 1.0) Ratio of dataset rows used while fitting each tree. max_depth : int or None (default = None) @@ -148,25 +148,25 @@ def fit(self, X, y, sample_weight=None): y : Dask cuDF DataFrame or CuPy backed Dask Array (n_rows, 1) Labels of training examples. **y must be partitioned the same way as X** - sample_weight : array-like, optional - Sample weights are not yet supported by distributed random - forests. + sample_weight : Dask cuDF Series or CuPy backed Dask Array, optional + Non-negative weights of training examples, partitioned the same + way as X and y. With ``bootstrap=True``, weights determine + sampling probabilities. Otherwise, weights are used in the + impurity calculation. """ - if self.kwargs.get("bootstrap", True) and sample_weight is not None: - raise NotImplementedError( - "Weighted bootstrapping is not yet supported for distributed " - "random forests. Set bootstrap=False or do not provide " - "sample_weight." - ) - if sample_weight is not None: - raise NotImplementedError( - "sample_weight is not yet supported for distributed random " - "forests" + if sample_weight is not None and ( + sample_weight.npartitions != X.npartitions + or sample_weight.npartitions != y.npartitions + ): + raise ValueError( + "X, y, and sample_weight must have the same number of " + "partitions" ) + self.internal_model = None self._fit( model=self.rfs, - dataset=(X, y), + dataset=(X, y) if sample_weight is None else (X, y, sample_weight), ) return self diff --git a/python/cuml/cuml/ensemble/randomforest_common.pyx b/python/cuml/cuml/ensemble/randomforest_common.pyx index 138443af5d..374f6cf149 100644 --- a/python/cuml/cuml/ensemble/randomforest_common.pyx +++ b/python/cuml/cuml/ensemble/randomforest_common.pyx @@ -431,6 +431,30 @@ class BaseRandomForestModel(InteropMixin, Base): stream=Stream.from_handle(stream), ) + def _fit(self, X, y, sample_weight=None): + """Prepare inputs once and coordinate validation across Dask workers.""" + validation_error = None + try: + X, y, sample_weight = self._prepare_fit_inputs(X, y, sample_weight) + if ( + hasattr(self, "_distributed_n_rows") + and not self.bootstrap + and sample_weight is not None + and sample_weight.sum().item() <= 0.0 + ): + raise ValueError( + "Rank-local sample weights must sum to a positive value" + ) + except Exception as error: + validation_error = error + + if self._allreduce_validation_status(validation_error is not None): + if validation_error is not None: + raise validation_error + raise RuntimeError("Input validation failed on another worker") + + return self._fit_forest(X, y, sample_weight=sample_weight) + def _fit_forest(self, X, y, sample_weight=None): cdef bool is_classifier = self._estimator_type == "classifier" cdef bool is_float32 = X.dtype == np.float32 diff --git a/python/cuml/cuml/ensemble/randomforestclassifier.py b/python/cuml/cuml/ensemble/randomforestclassifier.py index 4e91c98e8d..f6aff822b7 100644 --- a/python/cuml/cuml/ensemble/randomforestclassifier.py +++ b/python/cuml/cuml/ensemble/randomforestclassifier.py @@ -267,11 +267,14 @@ def fit(self, X, y, sample_weight=None) -> "RandomForestClassifier": """ Perform Random Forest Classification on the input data """ - X, y, sample_weight = self._prepare_fit_inputs(X, y, sample_weight) - return self._fit_forest(X, y, sample_weight=sample_weight) + return self._fit(X, y, sample_weight=sample_weight) def _prepare_fit_inputs(self, X, y, sample_weight=None): classes = getattr(self, "_distributed_classes", True) + # In distributed setting, some partitions may have sample weights that are all zeros + sample_weight_allow_all_zero = ( + hasattr(self, "_distributed_n_rows") and self.bootstrap + ) X, y, sample_weight, classes = check_inputs( self, X, @@ -281,6 +284,7 @@ def _prepare_fit_inputs(self, X, y, sample_weight=None): order="A", y_dtype="int32", sample_weight_dtype="float64", + sample_weight_allow_all_zero=sample_weight_allow_all_zero, return_classes=classes, reset=True, ) diff --git a/python/cuml/cuml/ensemble/randomforestregressor.py b/python/cuml/cuml/ensemble/randomforestregressor.py index f0fac0b95a..a658504907 100644 --- a/python/cuml/cuml/ensemble/randomforestregressor.py +++ b/python/cuml/cuml/ensemble/randomforestregressor.py @@ -211,10 +211,13 @@ def fit(self, X, y, sample_weight=None) -> "RandomForestRegressor": Perform Random Forest Regression on the input data """ - X, y, sample_weight = self._prepare_fit_inputs(X, y, sample_weight) - return self._fit_forest(X, y, sample_weight=sample_weight) + return self._fit(X, y, sample_weight=sample_weight) def _prepare_fit_inputs(self, X, y, sample_weight=None): + # In distributed setting, some partitions may have sample weights that are all zeros + sample_weight_allow_all_zero = ( + hasattr(self, "_distributed_n_rows") and self.bootstrap + ) return check_inputs( self, X, @@ -223,6 +226,7 @@ def _prepare_fit_inputs(self, X, y, sample_weight=None): dtype=("float32", "float64"), order="A", sample_weight_dtype="float64", + sample_weight_allow_all_zero=sample_weight_allow_all_zero, reset=True, ) diff --git a/python/cuml/cuml/internals/validation.py b/python/cuml/cuml/internals/validation.py index fdeae961d2..4bb40de118 100644 --- a/python/cuml/cuml/internals/validation.py +++ b/python/cuml/cuml/internals/validation.py @@ -1370,6 +1370,7 @@ def check_sample_weight( mem_type="device", order="A", ensure_non_negative=False, + allow_all_zero=False, ): """Validate and coerce ``sample_weight`` to a supported type. @@ -1395,6 +1396,9 @@ def check_sample_weight( ensure_non_negative : bool, default=False If True, an error will be raised if negative values are found in the input. By default ``check_non_negative`` is skipped. + allow_all_zero : bool, default=False + If True, allow ``sample_weight`` that are all zeros. + If False, all-zero ``sample_weight`` will be considered an error. Returns ------- @@ -1435,7 +1439,11 @@ def check_sample_weight( f"{sample_weight.ndim}D array." ) - if sample_weight.size and (sample_weight == 0).all(): + if ( + not allow_all_zero + and sample_weight.size + and (sample_weight == 0).all() + ): raise ValueError(all_zero_msg) return sample_weight @@ -1451,6 +1459,7 @@ def check_inputs( dtype=None, y_dtype=..., sample_weight_dtype=..., + sample_weight_allow_all_zero=False, mem_type="device", order="A", copy=False, @@ -1509,6 +1518,9 @@ def check_inputs( sample_weight_dtype : None, dtype, list[dtype], default=... The dtype(s) to support for sample_weight. If not specified, defaults to the output dtype of ``X``. + sample_weight_allow_all_zero : bool, default=False + If True, allow ``sample_weight`` that are all zeros. + If False, all-zero ``sample_weight`` will be considered an error. mem_type : {'device', 'host'} or None, default='device' The memory type use for the output. If 'device', the output will be a ``cupy.ndarray`` if dense, or a ``cupyx.scipy.sparse.spmatrix`` if @@ -1623,6 +1635,7 @@ def check_inputs( sample_weight = check_sample_weight( sample_weight, dtype=sample_weight_dtype, + allow_all_zero=sample_weight_allow_all_zero, mem_type=mem_type, order=order, ) diff --git a/python/cuml/cuml_accel_tests/test_onnx.py b/python/cuml/cuml_accel_tests/test_onnx.py index 368c12a25b..1c9dc0915c 100644 --- a/python/cuml/cuml_accel_tests/test_onnx.py +++ b/python/cuml/cuml_accel_tests/test_onnx.py @@ -65,7 +65,7 @@ def regression_data(): classifiers = [ pytest.param( - RandomForestClassifier(n_estimators=20, max_depth=8, random_state=42), + RandomForestClassifier(n_estimators=21, max_depth=8, random_state=42), marks=(xfail_skl2onnx_bug if ORT_VERSION < Version("1.28.0") else ()), id="RandomForestClassifier", ), diff --git a/python/cuml/cuml_accel_tests/upstream/scikit-learn/xfail-list.yaml b/python/cuml/cuml_accel_tests/upstream/scikit-learn/xfail-list.yaml index b05abe0192..211a512f08 100644 --- a/python/cuml/cuml_accel_tests/upstream/scikit-learn/xfail-list.yaml +++ b/python/cuml/cuml_accel_tests/upstream/scikit-learn/xfail-list.yaml @@ -273,8 +273,6 @@ - "sklearn.ensemble.tests.test_forest::test_max_samples_bootstrap[RandomForestClassifier]" - "sklearn.ensemble.tests.test_forest::test_max_samples_bootstrap[RandomForestRegressor]" - "sklearn.ensemble.tests.test_forest::test_min_impurity_decrease" - - "sklearn.ensemble.tests.test_forest::test_min_samples_leaf[RandomForestClassifier]" - - "sklearn.ensemble.tests.test_forest::test_min_samples_leaf[RandomForestRegressor]" - "sklearn.ensemble.tests.test_forest::test_oob_not_computed_twice[RandomForestClassifier]" - "sklearn.ensemble.tests.test_forest::test_oob_not_computed_twice[RandomForestRegressor]" - "sklearn.ensemble.tests.test_forest::test_poisson_y_positive_check" @@ -450,7 +448,6 @@ - "sklearn.tests.test_common::test_estimators[Pipeline(steps=[('scaler',StandardScaler()),('final_estimator',LogisticRegression())])-check_estimators_overwrite_params]" - "sklearn.tests.test_common::test_estimators[Pipeline(steps=[('scaler',StandardScaler()),('final_estimator',Ridge())])-check_dont_overwrite_parameters]" - "sklearn.tests.test_common::test_estimators[Pipeline(steps=[('scaler',StandardScaler()),('final_estimator',Ridge())])-check_estimators_overwrite_params]" - - "sklearn.tests.test_multioutput::test_multi_output_classification_sample_weights" - "sklearn.tests.test_multioutput::test_multi_target_sparse_regression[bsr_array]" - "sklearn.tests.test_multioutput::test_multi_target_sparse_regression[bsr_matrix]" - "sklearn.tests.test_multioutput::test_multi_target_sparse_regression[coo_array]" diff --git a/python/cuml/tests/dask/test_dask_random_forest.py b/python/cuml/tests/dask/test_dask_random_forest.py index d40368ad8c..aa8c31c3b0 100644 --- a/python/cuml/tests/dask/test_dask_random_forest.py +++ b/python/cuml/tests/dask/test_dask_random_forest.py @@ -451,94 +451,214 @@ def distributed_data(y_parts): ) -def test_rf_classification_zero_class_weight_on_one_worker(client): +@pytest.mark.parametrize("bootstrap", [False, True]) +@pytest.mark.parametrize("weight_type", ["sample", "class"]) +def test_rf_zero_weight_on_one_worker(client, bootstrap, weight_type): workers = list(client.scheduler_info(n_workers=-1)["workers"])[:2] if len(workers) < 2: pytest.skip("This test requires at least two workers") X_futures = [] y_futures = [] + weight_futures = [] for label, worker in enumerate(workers): X_part = cudf.DataFrame(np.zeros((20, 1), dtype=np.float32)) y_part = cudf.Series(np.full(20, label, dtype=np.int32)) X_futures.append(client.scatter(X_part, workers=[worker], hash=False)) y_futures.append(client.scatter(y_part, workers=[worker], hash=False)) + weight_part = cudf.Series(np.full(20, label, dtype=np.float64)) + weight_futures.append( + client.scatter(weight_part, workers=[worker], hash=False) + ) X = dask_cudf.from_delayed(X_futures, meta=X_part.iloc[:0]) y = dask_cudf.from_delayed(y_futures, meta=y_part.iloc[:0]) + weights = dask_cudf.from_delayed(weight_futures, meta=weight_part.iloc[:0]) + kwargs = ( + {"class_weight": {0: 0.0, 1: 1.0}} if weight_type == "class" else {} + ) + fit_kwargs = {"sample_weight": weights} if weight_type == "sample" else {} model = cuRFC_mg( workers=workers, n_estimators=1, - bootstrap=False, + bootstrap=bootstrap, max_depth=1, n_bins=2, random_state=42, - class_weight={0: 0.0, 1: 1.0}, + **kwargs, ) + if bootstrap: + with dask.annotate(workers=workers): + model.fit(X, y, **fit_kwargs) + combined = model.get_combined_model() + combined.output_type = "cupy" + prediction = combined.predict(cp.zeros((1, 1), dtype=cp.float32)) + cp.testing.assert_array_equal(prediction, [1]) + return + with pytest.raises( RuntimeError, match="2 of 2 worker jobs failed" ) as exc_info: with dask.annotate(workers=workers): - model.fit(X, y) + model.fit(X, y, **fit_kwargs) - assert "Rank-local sample weights must sum to a positive value" in str( - exc_info.value + message = ( + "Rank-local sample weights must sum to a positive value" + if weight_type == "class" + else "Sample weights must contain at least one non-zero number" ) + assert message in str(exc_info.value) +@pytest.mark.parametrize("bootstrap", [False, True]) +@pytest.mark.parametrize("input_type", ["dataframe", "array"]) @pytest.mark.parametrize( - "model_cls,model_kwargs,fit_kwargs", + "model_cls,class_weight,use_sample_weight,expected", [ - pytest.param( - cuRFC_mg, - {"class_weight": "balanced"}, - {}, - id="classifier-class-weight", - ), - pytest.param( - cuRFC_mg, - {}, - {"sample_weight": np.ones(4, dtype=np.float64)}, - id="classifier-sample-weight", - ), - pytest.param( - cuRFR_mg, - {}, - {"sample_weight": np.ones(4, dtype=np.float64)}, - id="regressor-sample-weight", - ), - pytest.param( - cuRFC_mg, - {"class_weight": "balanced"}, - {"sample_weight": np.ones(4, dtype=np.float64)}, - id="classifier-sample-and-class-weight", - ), + pytest.param(cuRFC_mg, "balanced", False, 0.5, id="balanced"), + pytest.param(cuRFC_mg, {0: 9, 1: 1}, False, 0.25, id="class"), + pytest.param(cuRFC_mg, None, True, 0.9, id="sample-classifier"), + pytest.param(cuRFR_mg, None, True, 0.9, id="sample-regressor"), + pytest.param(cuRFC_mg, {0: 9, 1: 1}, True, 0.5, id="combined"), + pytest.param(cuRFC_mg, "balanced", True, 0.5, id="balanced-combined"), ], ) -def test_rf_weighted_bootstrap_raises( - client, model_cls, model_kwargs, fit_kwargs +def test_rf_weights( + client, + bootstrap, + input_type, + model_cls, + class_weight, + use_sample_weight, + expected, ): - X = dask_cudf.from_cudf( - cudf.DataFrame(np.arange(8, dtype=np.float32).reshape(4, 2)), - npartitions=1, - ) + # Constant features make leaf values equal to the weighted label mean. + # The imbalanced labels distinguish global balancing from local balancing. + n_workers = len(client.scheduler_info(n_workers=-1)["workers"]) + n_rows = 400 * n_workers + X = np.zeros((n_rows, 1), dtype=np.float32) y_dtype = np.int32 if model_cls is cuRFC_mg else np.float32 - y = dask_cudf.from_cudf( - cudf.Series(np.arange(4, dtype=y_dtype) % 2), npartitions=1 - ) + y = np.repeat([0, 1], [n_rows // 4, 3 * n_rows // 4]).astype(y_dtype) + weights = np.where(y == 0, 1.0, 3.0) + if input_type == "dataframe": + data = [ + dask_cudf.from_cudf(cudf.DataFrame(X), npartitions=n_workers * 2), + dask_cudf.from_cudf(cudf.Series(y), npartitions=n_workers * 2), + dask_cudf.from_cudf( + cudf.Series(weights), npartitions=n_workers * 2 + ), + ] + else: + data = [from_array(cp.asarray(a), chunks=200) for a in (X, y, weights)] + X, y, weights = dask_utils.persist_across_workers(client, data) + kwargs = {} if class_weight is None else {"class_weight": class_weight} model = model_cls( + n_estimators=16, + bootstrap=bootstrap, + max_depth=1, + n_bins=2, + random_state=42, + **kwargs, + ) + initial_models = model.rfs.copy() + model.fit(X, y, sample_weight=weights if use_sample_weight else None) + + combined = model.get_combined_model() + combined.output_type = "cupy" + X_test = cp.zeros((1, 1), dtype=cp.float32) + if model_cls is cuRFC_mg: + prediction = combined.predict_proba(X_test)[0, 1] + else: + prediction = combined.predict(X_test)[0] + assert float(prediction) == pytest.approx( + expected, abs=0.03 if bootstrap else 1e-6 + ) + + # All participating ranks must build the same forest. Workers without + # training rows keep their initial, unfitted model. + worker_model_bytes = client.gather( + [ + client.submit(_get_treelite_bytes, model.rfs[w], workers=[w]) + for w in model.rfs + if model.rfs[w] is not initial_models[w] + ] + ) + assert all(data == worker_model_bytes[0] for data in worker_model_bytes) + + +def test_rf_with_column_shaped_label(client): + n_workers = len(client.scheduler_info(n_workers=-1)["workers"]) + n_rows = 400 * n_workers + X = cp.zeros((n_rows, 1), dtype=cp.float32) + y = cp.repeat( + cp.asarray([0, 1], dtype=cp.int32), + [n_rows // 4, 3 * n_rows // 4], + ).reshape(-1, 1) + sample_weight = cp.where(y.ravel() == 0, 1.0, 3.0) + + X, y, sample_weight = dask_utils.persist_across_workers( + client, + [from_array(a, chunks=200) for a in (X, y, sample_weight)], + ) + model = cuRFC_mg( n_estimators=1, - bootstrap=True, + bootstrap=False, max_depth=1, n_bins=2, - **model_kwargs, + random_state=42, + class_weight="balanced", + ).fit(X, y, sample_weight=sample_weight) + + combined = model.get_combined_model() + combined.output_type = "cupy" + prediction = combined.predict_proba(cp.zeros((1, 1), dtype=cp.float32))[ + 0, 1 + ] + assert float(prediction) == pytest.approx(0.5, abs=1e-6) + + +@pytest.mark.parametrize("model_cls", [cuRFC_mg, cuRFR_mg]) +def test_rf_sample_weight_partition_count_mismatch(client, model_cls): + X = dask_cudf.from_cudf( + cudf.DataFrame(np.zeros((20, 1), dtype=np.float32)), npartitions=2 + ) + y_dtype = np.int32 if model_cls is cuRFC_mg else np.float32 + y = dask_cudf.from_cudf( + cudf.Series(np.arange(20, dtype=y_dtype) % 2), npartitions=2 + ) + sample_weight = dask_cudf.from_cudf( + cudf.Series(np.ones(10, dtype=np.float32)), npartitions=1 ) + model = model_cls(n_estimators=1) + with pytest.raises( + ValueError, + match="X, y, and sample_weight must have the same number of partitions", + ): + model.fit(X, y, sample_weight=sample_weight) + + +@pytest.mark.parametrize("weight_type", ["sample", "class"]) +def test_rf_weighted_bootstrap_all_zero_weights(client, weight_type): + n_workers = len(client.scheduler_info(n_workers=-1)["workers"]) + X, y = _prep_training_data( + client, + np.zeros((n_workers * 20, 1), dtype=np.float32), + np.arange(n_workers * 20, dtype=np.int32) % 2, + partitions_per_worker=1, + ) + kwargs = ( + {"class_weight": {0: 0.0, 1: 0.0}} if weight_type == "class" else {} + ) + fit_kwargs = {"sample_weight": y * 0.0} if weight_type == "sample" else {} + model = cuRFC_mg( + n_estimators=1, bootstrap=True, max_depth=1, n_bins=2, **kwargs + ) with pytest.raises( - NotImplementedError, - match="Weighted bootstrapping is not yet supported", + RuntimeError, + match="sample_weight values must contain at least one positive value", ): model.fit(X, y, **fit_kwargs) diff --git a/python/cuml/tests/test_validation.py b/python/cuml/tests/test_validation.py index e607521220..d5ac9d0b4f 100644 --- a/python/cuml/tests/test_validation.py +++ b/python/cuml/tests/test_validation.py @@ -1869,6 +1869,13 @@ def test_check_sample_weight_all_zero(sample_weight): check_sample_weight(sample_weight) +@pytest.mark.parametrize("mem_type", ["host", "device"]) +def test_check_sample_weight_allow_all_zero(mem_type): + weights = np.zeros(3) + out = check_sample_weight(weights, mem_type=mem_type, allow_all_zero=True) + np.testing.assert_array_equal(cp.asnumpy(out), weights) + + @pytest.mark.parametrize("value", ["NaN", "infinity"]) def test_check_sample_weight_non_finite(value): scalar = float(value)