Add class_weight='balanced_subsample' to RandomForest and ExtraTrees - #8181
Closed
switch527 wants to merge 3 commits into
Closed
Add class_weight='balanced_subsample' to RandomForest and ExtraTrees#8181switch527 wants to merge 3 commits into
switch527 wants to merge 3 commits into
Conversation
Thread true per-sample `sample_weight` through the RandomForest C++ tree builder and add classifier `class_weight`, bringing `RandomForestClassifier` and `RandomForestRegressor` in line with scikit-learn's weighting API. Weights enter split finding and leaf statistics directly. The unweighted path stays byte-identical to the prior behavior; the weighted path uses a deterministic pool-based cross-block reduction.
Closes NVIDIA#8133. Adds `ExtraTreesClassifier` and `ExtraTreesRegressor` to `cuml.ensemble`. Both subclass the existing RandomForest estimators, route through a new `SPLITTER_RANDOM` C++ path, and default to `bootstrap=False` to match scikit-learn's ExtraTrees defaults. `sample_weight` and `class_weight` reuse the parent estimators' machinery. The C++ side is a direct-single-threshold split kernel that skips the per-bin histogram. Splitter-aware dispatch in `Builder::computeSplit` keeps the `SPLITTER_BEST` arm byte-identical; the `SPLITTER_RANDOM` arm uses a smaller per-block scratch and reuses the deterministic pool primitive introduced by NVIDIA#8132. Adjacent fix: `BaseRandomForestModel._params_to_cpu` coerces `max_samples` to `None` when `bootstrap=False` so the sklearn round-trip succeeds. Affects RandomForest too; sklearn rejects that combination anyway. Test coverage adds a new C++ gtest binary plus Python sklearn-parity, round-trip, and accel-integration tests for both estimators. Stacks on NVIDIA#8132 since the weighted bin types and pool primitive live there. The Python entry point routes through the existing RandomForest code path unchanged; `_splitter` is a class attribute, not a user-facing parameter.
Closes NVIDIA#8146. Computes per-tree class weights on-device from each bootstrap's class distribution and folds them into the split-finding and leaf-prediction histogram kernels under a HasTreeClassWeight template parameter. bootstrap=False collapses to 'balanced' via a private _effective_class_weight property and emits a UserWarning; self.class_weight stays untouched so clone() round-trips.
switch527
force-pushed
the
fea-rf-balanced-subsample
branch
from
May 29, 2026 03:49
048e3fa to
44013e1
Compare
Contributor
Author
📝 WalkthroughSummary by CodeRabbitRelease Notes
WalkthroughImplements ExtraTrees (random splitter) and end-to-end weighted training: new enums/params, weighted bins/objectives, builder/kernel updates (including random-split kernels), RF integration with per-tree class weights, Python bindings and new ExtraTrees estimators, plus comprehensive C++/Python tests and benchmarks. Build/CMake updated to compile new kernels. ChangesCore Tree/Forest Contracts
Weighted Bins and Objectives
Builder and Kernels
RandomForest Integration
Python Bindings and ExtraTrees
C++ Tests
Python Tests, Dask, and Compatibility
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related issues
Possibly related PRs
Suggested labels
Suggested reviewers
✨ Finishing Touches🧪 Generate unit tests (beta)
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
python/cuml/cuml_accel_tests/integration/test_et_classifier.py (1)
46-53: ⚡ Quick winAdd explicit
balanced_subsamplecoverage in class_weight parametrization.Line 46 currently skips the new mode introduced by this PR, so the integration suite can miss regressions on the core behavior.
✅ Suggested test update
-@pytest.mark.parametrize("class_weight", [None, "balanced", {0: 1, 1: 2}]) +@pytest.mark.parametrize( + "class_weight", + [None, "balanced", "balanced_subsample", {0: 1, 1: 2}], +) def test_et_class_weight(classification_data, class_weight):As per coding guidelines
**/*test*.{py,cpp,hpp}: Update unit tests when implementing code changes.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cuml/cuml_accel_tests/integration/test_et_classifier.py` around lines 46 - 53, The test test_et_class_weight for ExtraTreesClassifier doesn't include the new "balanced_subsample" mode; update the pytest.mark.parametrize on class_weight in test_et_class_weight to include "balanced_subsample" alongside None, "balanced", and the dict {0:1, 1:2} so the ExtraTreesClassifier class_weight behavior (constructor/fit/predict) is exercised for the new mode and the integration suite catches regressions.cpp/src/decisiontree/batched-levelalgo/builder.cuh (1)
21-31: ⚡ Quick winInclude
<tuple>explicitly for the newstd::make_tupleusage.
updateWorkloadInfo()anddoSplit()now depend onstd::make_tuple(), but this header still builds only if some transitive include happens to pull in<tuple>. Please add the direct include here so the file stays self-contained.Suggested fix
`#include` <deque> `#include` <memory> +#include <tuple> `#include` <type_traits> `#include` <utility>🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/decisiontree/batched-levelalgo/builder.cuh` around lines 21 - 31, The file is missing a direct include of <tuple>, causing builds to rely on transitive includes for new uses of std::make_tuple in updateWorkloadInfo and doSplit; add `#include` <tuple> at the top of builder.cuh (near the other STL includes) so the file is self-contained and both updateWorkloadInfo and doSplit can use std::make_tuple reliably.cpp/src/decisiontree/batched-levelalgo/objectives.cuh (1)
753-762: 💤 Low valuePotential division by zero in
SetLeafVector.If
totalsums to zero (allweighted_sumvalues are zero due to edge cases or floating-point accumulation), this division produces undefined behavior. Per coding guidelines, add an epsilon check.This pattern appears in all six weighted objective
SetLeafVectormethods (lines 753-762, 871-879, 961-966, 1054-1058, 1147-1151, 1238-1242).🛡️ Suggested epsilon guard
static DI void SetLeafVector(BinT const* shist, int nclasses, DataT* out) { double total = 0.0; for (int i = 0; i < nclasses; i++) { total += shist[i].weighted_sum; } + // Guard against degenerate all-zero-weight leaf + if (total <= 0.0) total = 1.0; for (int i = 0; i < nclasses; i++) { out[i] = DataT(shist[i].weighted_sum) / DataT(total); } }As per coding guidelines: "Add epsilon checks for division by zero or near-zero values".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/decisiontree/batched-levelalgo/objectives.cuh` around lines 753 - 762, The SetLeafVector function may divide by zero when total (sum of shist[i].weighted_sum) is zero; update SetLeafVector (and the other five weighted SetLeafVector overloads) to compute total, compare it against a small epsilon (e.g., 1e-12) and, if total <= epsilon, populate out[] with a safe default (suggest using uniform probabilities DataT(1)/DataT(nclasses)), otherwise proceed with out[i] = DataT(shist[i].weighted_sum) / DataT(total); reference symbols: SetLeafVector, BinT, DataT, shist, nclasses, out.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/cuml/cuml/ensemble/randomforest_common.pyx`:
- Around line 494-500: Validate sample_weight before taking its pointer: ensure
sample_weight (the input used with cp.ascontiguousarray and assigned to
sample_weight_dev/sample_weight_ptr) is a 1-D array with length equal to the
number of rows in X (e.g., X.shape[0]); if not, raise a ValueError describing
the mismatch. Convert to a CuPy array first (cp.asarray), check ndim == 1 and
size == X.shape[0], then call cp.ascontiguousarray(..., dtype=X.dtype) and take
sample_weight_dev.data.ptr only after those checks pass.
---
Nitpick comments:
In `@cpp/src/decisiontree/batched-levelalgo/builder.cuh`:
- Around line 21-31: The file is missing a direct include of <tuple>, causing
builds to rely on transitive includes for new uses of std::make_tuple in
updateWorkloadInfo and doSplit; add `#include` <tuple> at the top of builder.cuh
(near the other STL includes) so the file is self-contained and both
updateWorkloadInfo and doSplit can use std::make_tuple reliably.
In `@cpp/src/decisiontree/batched-levelalgo/objectives.cuh`:
- Around line 753-762: The SetLeafVector function may divide by zero when total
(sum of shist[i].weighted_sum) is zero; update SetLeafVector (and the other five
weighted SetLeafVector overloads) to compute total, compare it against a small
epsilon (e.g., 1e-12) and, if total <= epsilon, populate out[] with a safe
default (suggest using uniform probabilities DataT(1)/DataT(nclasses)),
otherwise proceed with out[i] = DataT(shist[i].weighted_sum) / DataT(total);
reference symbols: SetLeafVector, BinT, DataT, shist, nclasses, out.
In `@python/cuml/cuml_accel_tests/integration/test_et_classifier.py`:
- Around line 46-53: The test test_et_class_weight for ExtraTreesClassifier
doesn't include the new "balanced_subsample" mode; update the
pytest.mark.parametrize on class_weight in test_et_class_weight to include
"balanced_subsample" alongside None, "balanced", and the dict {0:1, 1:2} so the
ExtraTreesClassifier class_weight behavior (constructor/fit/predict) is
exercised for the new mode and the integration suite catches regressions.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 603930db-0c4b-460c-8f0f-a8d9440a71c3
📒 Files selected for processing (77)
cpp/CMakeLists.txtcpp/include/cuml/ensemble/randomforest.hppcpp/include/cuml/tree/decisiontree.hppcpp/src/decisiontree/batched-levelalgo/bins.cuhcpp/src/decisiontree/batched-levelalgo/builder.cuhcpp/src/decisiontree/batched-levelalgo/dataset.hcpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuhcpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels_impl.cuhcpp/src/decisiontree/batched-levelalgo/kernels/builder_random_kernels_impl.cuhcpp/src/decisiontree/batched-levelalgo/kernels/random_entropy-double.cucpp/src/decisiontree/batched-levelalgo/kernels/random_entropy-float.cucpp/src/decisiontree/batched-levelalgo/kernels/random_gamma-double.cucpp/src/decisiontree/batched-levelalgo/kernels/random_gamma-float.cucpp/src/decisiontree/batched-levelalgo/kernels/random_gini-double.cucpp/src/decisiontree/batched-levelalgo/kernels/random_gini-float.cucpp/src/decisiontree/batched-levelalgo/kernels/random_inverse_gaussian-double.cucpp/src/decisiontree/batched-levelalgo/kernels/random_inverse_gaussian-float.cucpp/src/decisiontree/batched-levelalgo/kernels/random_mse-double.cucpp/src/decisiontree/batched-levelalgo/kernels/random_mse-float.cucpp/src/decisiontree/batched-levelalgo/kernels/random_poisson-double.cucpp/src/decisiontree/batched-levelalgo/kernels/random_poisson-float.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted_entropy-double.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted_entropy-float.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted_gamma-double.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted_gamma-float.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted_gini-double.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted_gini-float.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted_inverse_gaussian-double.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted_inverse_gaussian-float.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted_mse-double.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted_mse-float.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted_poisson-double.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted_poisson-float.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted_random_entropy-double.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted_random_entropy-float.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted_random_gamma-double.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted_random_gamma-float.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted_random_gini-double.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted_random_gini-float.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted_random_inverse_gaussian-double.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted_random_inverse_gaussian-float.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted_random_mse-double.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted_random_mse-float.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted_random_poisson-double.cucpp/src/decisiontree/batched-levelalgo/kernels/weighted_random_poisson-float.cucpp/src/decisiontree/batched-levelalgo/objectives.cuhcpp/src/decisiontree/decisiontree.cucpp/src/decisiontree/decisiontree.cuhcpp/src/randomforest/randomforest.cucpp/src/randomforest/randomforest.cuhcpp/tests/CMakeLists.txtcpp/tests/sg/extratrees_rng_reference.hcpp/tests/sg/extratrees_rng_reference_gen.pycpp/tests/sg/extratrees_test.cucpp/tests/sg/rf_test.cupython/cuml/cuml/__init__.pypython/cuml/cuml/accel/_overrides/sklearn/ensemble.pypython/cuml/cuml/benchmark/algorithms.pypython/cuml/cuml/benchmark/automated/bench_extra_trees.pypython/cuml/cuml/dask/ensemble/randomforestclassifier.pypython/cuml/cuml/dask/ensemble/randomforestregressor.pypython/cuml/cuml/ensemble/__init__.pypython/cuml/cuml/ensemble/extra_trees_classifier.pypython/cuml/cuml/ensemble/extra_trees_regressor.pypython/cuml/cuml/ensemble/randomforest_common.pyxpython/cuml/cuml/ensemble/randomforestclassifier.pypython/cuml/cuml/ensemble/randomforestregressor.pypython/cuml/cuml_accel_tests/integration/test_et_classifier.pypython/cuml/cuml_accel_tests/integration/test_et_regressor.pypython/cuml/cuml_accel_tests/integration/test_rf_classifier.pypython/cuml/cuml_accel_tests/integration/test_rf_regressor.pypython/cuml/cuml_accel_tests/upstream/scikit-learn/xfail-list.yamlpython/cuml/tests/dask/test_dask_random_forest.pypython/cuml/tests/test_extratrees.pypython/cuml/tests/test_random_forest.pypython/cuml/tests/test_sklearn_compatibility.pypython/cuml/tests/test_sklearn_import_export.py
Comment on lines
+494
to
+500
| # None -> NULL pointer -> unweighted C++ path unchanged. | ||
| sample_weight_dev = None | ||
| cdef uintptr_t sample_weight_ptr = 0 | ||
| if sample_weight is not None: | ||
| sample_weight_dev = cp.ascontiguousarray( | ||
| sample_weight, dtype=X.dtype) | ||
| sample_weight_ptr = sample_weight_dev.data.ptr |
There was a problem hiding this comment.
Validate sample_weight before taking its raw pointer.
cp.ascontiguousarray() will accept a mis-sized or multi-dimensional input, and the C++ fit path will still read it as n_rows contiguous weights. That can silently train on the wrong weights or walk past the intended buffer. Please enforce a 1-D array with exactly one entry per row here.
Suggested fix
sample_weight_dev = None
cdef uintptr_t sample_weight_ptr = 0
if sample_weight is not None:
+ if sample_weight.ndim != 1 or sample_weight.shape[0] != X.shape[0]:
+ raise ValueError(
+ "sample_weight must be 1-dimensional with one entry per row"
+ )
sample_weight_dev = cp.ascontiguousarray(
sample_weight, dtype=X.dtype)
sample_weight_ptr = sample_weight_dev.data.ptr🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cuml/cuml/ensemble/randomforest_common.pyx` around lines 494 - 500,
Validate sample_weight before taking its pointer: ensure sample_weight (the
input used with cp.ascontiguousarray and assigned to
sample_weight_dev/sample_weight_ptr) is a 1-D array with length equal to the
number of rows in X (e.g., X.shape[0]); if not, raise a ValueError describing
the mismatch. Convert to a CuPy array first (cp.asarray), check ndim == 1 and
size == X.shape[0], then call cp.ascontiguousarray(..., dtype=X.dtype) and take
sample_weight_dev.data.ptr only after those checks pass.
chyunsu3
added
improvement
switch527
marked this pull request as draft
May 30, 2026 16:56
Contributor
Author
|
Covered by the new PR structure for class weight and extra trees |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #8146. Stacked on #8132 (RF
sample_weight+ classifierclass_weight) and #8143 (ExtraTrees estimators). Review the third commit for the balanced_subsample-specific surface; the first two commits are the dependencies and will collapse on merge.Adds
class_weight='balanced_subsample'toRandomForestClassifierandExtraTreesClassifier, completing sklearnclass_weightparity. The per-tree compute and theHasTreeClassWeightkernel template gating only make sense on top of #8132 + #8143.The mode computes class weights on-device per tree from the bootstrap's class distribution and folds them into the histogram kernel's per-row weight under a compile-time
HasTreeClassWeighttemplate parameter. The split-finding kernel and the leaf-prediction kernel both apply the multiplier; without the leaf-side fix, leaf probabilities collapse to the unweighted majority and minority recall goes the wrong way (caught by the Phase 6 minority-recall test). 12 new explicit instantiations gated to weighted-classifier.cufiles via a#define INSTANTIATE_TREE_CLASS_WEIGHTblock in the impl headers; weighted-regression files don't define the macro sotree_class_weight[floating-point label]never tries to compile.bootstrap=Falsehas no per-tree bootstrap, so a private@property _effective_class_weightcollapses the mode to'balanced'andfit()emits aUserWarningonce. The property is the source of truth for both Python and Cython reads;self.class_weightstays untouched soclone()round-trips. sklearn falls back silently here; cuml is explicit per the property pattern.feature_importances_underbalanced_subsampledoes not match sklearn's formula shape: cuml weights the importance at the impurity term via the kernel multiplier, sklearn weights it at the count term viaweighted_n_node_samples. Both apply the per-tree class weight once. Closing the formula gap would require threadingtree_class_weightinto the node-count kernels and is out of scope here.Three pre-PR tests that asserted
balanced_subsampleraises (test_rf_classifier_class_weight_balanced_subsample_raises[*],test_class_weight_balanced_subsample_rejected) are replaced by positive forms: minority-recall lift on a 90:10 imbalanced fixture, bootstrap=False collapse to'balanced'viacp.testing.assert_array_equalon predictions, sklearnfrom_sklearnround-trip, and a feature_importances well-formedness check. Thebootstrap=FalseUserWarning has paired positive + negative tests.test_sklearn_import_export.pyparametrize lists for RFC and ETC now include'balanced_subsample'.Binary size delta: +97 KB on libcuml.so for the 12 new instantiations and the per-tree compute kernel.