Thread sample weights through RandomForest Python API

cpp/src/randomforest/randomforest.cuh

@@ -20,9 +20,11 @@
 #include <rmm/exec_policy.hpp>
 
 #include <thrust/binary_search.h>
+#include <thrust/copy.h>
 #include <thrust/fill.h>
 #include <thrust/for_each.h>
 #include <thrust/iterator/constant_iterator.h>
+#include <thrust/iterator/counting_iterator.h>
 #include <thrust/logical.h>
 #include <thrust/reduce.h>
 #include <thrust/scan.h>
@@ -51,6 +53,11 @@ struct InvalidSampleWeight {
   __device__ bool operator()(T weight) const { return weight < T(0) || !isfinite(weight); }
 };
 
+template <typename T>
+struct NonzeroSampleWeight {
+  __device__ bool operator()(T weight) const { return weight != T(0); }
+};
+
 // Matches estimator behavior: when bootstrapping is enabled and sample weights exist,
 // those weights are materialized by drawing bootstrap rows according to them.
 class RowSampler {
@@ -106,34 +113,49 @@ class RowSampler {
 
     auto& selected_rows = selected_rows_[stream_id];
 
+    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);
     raft::random::RngState rng_state(rs, raft::random::GenPhilox);
 
-    if (bootstrap_) {
-      raft::resources stream_resources;
-      raft::resource::set_cuda_stream(stream_resources, stream);
-      if (use_weighted_bootstrap()) {
-        auto& weighted_draw_scratch = weighted_draw_scratch_[stream_id];
-        raft::random::uniform<double>(stream_resources,
-                                      rng_state,
-                                      weighted_draw_scratch.data(),
-                                      weighted_draw_scratch.size(),
-                                      0.0,
-                                      sample_weight_sum_);
-        thrust::upper_bound(rmm::exec_policy(stream),
-                            sample_weight_cdf_.data(),
-                            sample_weight_cdf_.data() + n_rows_,
-                            weighted_draw_scratch.begin(),
-                            weighted_draw_scratch.end(),
-                            selected_rows.begin());
-      } else {
-        raft::random::uniformInt<int>(
-          stream_resources, rng_state, selected_rows.data(), selected_rows.size(), 0, n_rows_);
-      }
+    if (use_weighted_bootstrap()) {
+      // Draw bootstrap rows according to sample weights.
+      auto& weighted_draw_scratch = weighted_draw_scratch_[stream_id];
+      raft::random::uniform<double>(stream_resources,
+                                    rng_state,
+                                    weighted_draw_scratch.data(),
+                                    weighted_draw_scratch.size(),
+                                    0.0,
+                                    sample_weight_sum_);
+      thrust::upper_bound(rmm::exec_policy(stream),
+                          sample_weight_cdf_.data(),
+                          sample_weight_cdf_.data() + n_rows_,
+                          weighted_draw_scratch.begin(),
+                          weighted_draw_scratch.end(),
+                          selected_rows.begin());
+    } else if (bootstrap_) {
+      // Draw bootstrap rows uniformly when there are no sample weights.
+      raft::random::uniformInt<int>(
+        stream_resources, rng_state, selected_rows.data(), selected_rows.size(), 0, n_rows_);
+    } else if (sample_weight_ != nullptr) {
+      // Remove zero-weight rows from the non-bootstrap row set.
+      selected_rows.resize(n_sampled_rows_, stream);
+      auto rows_begin        = thrust::make_counting_iterator<int>(0);
+      auto selected_rows_end = thrust::copy_if(rmm::exec_policy(stream),
+                                               rows_begin,
+                                               rows_begin + n_rows_,
+                                               sample_weight_,
+                                               selected_rows.begin(),
+                                               NonzeroSampleWeight<double>{});
+      auto n_selected        = selected_rows_end - selected_rows.begin();
+      ASSERT(n_selected > 0, "sample_weight values must contain at least one positive value");
+      selected_rows.resize(n_selected, stream);
     } else {
+      selected_rows.resize(n_sampled_rows_, stream);
       thrust::sequence(rmm::exec_policy(stream), selected_rows.begin(), selected_rows.end());
     }
 
@@ -155,7 +177,7 @@ class RowSampler {
     thrust::fill(rmm::exec_policy(stream), tree_mask, tree_mask + n_rows_, false);
     thrust::scatter(rmm::exec_policy(stream),
                     thrust::make_constant_iterator(true),
-                    thrust::make_constant_iterator(true) + n_sampled_rows_,
+                    thrust::make_constant_iterator(true) + selected_rows.size(),
                     selected_rows.data(),
                     tree_mask);
   }
@@ -256,7 +278,8 @@ class RandomForest {
   *   (n_trees * n_rows), only populated if a non-null pointer is provided.
   * @param[in] sample_weight: optional device pointer to per-row sample weights. With bootstrap
   *   enabled, rows are sampled with probability proportional to these weights and the sampled
-  *   counts drive tree training. Without bootstrap, weights are used for impurity/objective math.
+  *   counts drive tree training. Without bootstrap, zero-weight rows are removed from the tree
+  *   row set and remaining weights are used for impurity/objective math.
   */
   void fit(const raft::handle_t& user_handle,
            const T* input,
@@ -310,7 +333,7 @@ class RandomForest {
 
       /* Build individual tree in the forest.
         - input is a pointer to orig data that have n_cols features and n_rows rows.
-        - n_sampled_rows: # rows sampled for tree's bootstrap sample.
+        - n_sampled_rows: # rows sampled or retained for this tree.
         - sorted_selected_rows: points to a list of row #s (w/ n_sampled_rows elements)
           used to build the bootstrapped sample.
           Expectation: Each tree node will contain (a) # n_sampled_rows and

cpp/tests/sg/rf_test.cu

@@ -1643,7 +1643,7 @@ TEST(RfWeightedTest, RegressionRootLeafUsesWeights)
   const auto& tree = *forest->trees[0];
   ASSERT_EQ(tree.sparsetree.size(), 1);
   EXPECT_TRUE(tree.sparsetree[0].IsLeaf());
-  EXPECT_EQ(tree.sparsetree[0].InstanceCount(), 3);
+  EXPECT_EQ(tree.sparsetree[0].InstanceCount(), 2);
   ASSERT_EQ(tree.vector_leaf.size(), 1);
   EXPECT_NEAR(tree.vector_leaf[0], 7.5f, 1e-6f);
 }
@@ -1718,7 +1718,7 @@ TEST(RfWeightedTest, ZeroWeightSamplesDoNotCreatePositiveWeightSplit)
   const auto& tree = *forest->trees[0];
   ASSERT_EQ(tree.sparsetree.size(), 1);
   EXPECT_TRUE(tree.sparsetree[0].IsLeaf());
-  EXPECT_EQ(tree.sparsetree[0].InstanceCount(), 4);
+  EXPECT_EQ(tree.sparsetree[0].InstanceCount(), 2);
   ASSERT_EQ(tree.vector_leaf.size(), 2);
   EXPECT_NEAR(tree.vector_leaf[0], 0.0f, 1e-6f);
   EXPECT_NEAR(tree.vector_leaf[1], 1.0f, 1e-6f);

python/cuml/cuml/accel/_overrides/sklearn/ensemble.py

@@ -25,9 +25,6 @@ def _check_inputs(self, X, y=None, sample_weight=None):
                 ) from None
             raise
 
-        if sample_weight is not None:
-            raise UnsupportedOnGPU("`sample_weight` is not supported")
-
         if y is not None:
             y = check_array(
                 y,
@@ -48,15 +45,15 @@ def _check_inputs(self, X, y=None, sample_weight=None):
 
     def _gpu_fit(self, X, y, sample_weight=None):
         self._check_inputs(X, y, sample_weight=sample_weight)
-        return self._gpu.fit(X, y)
+        return self._gpu.fit(X, y, sample_weight=sample_weight)
 
     def _gpu_predict(self, X):
         self._check_inputs(X)
         return self._gpu.predict(X)
 
     def _gpu_score(self, X, y, sample_weight=None):
         self._check_inputs(X, y, sample_weight=sample_weight)
-        return self._gpu.score(X, y)
+        return self._gpu.score(X, y, sample_weight=sample_weight)
 
 
 class RandomForestRegressor(ProxyBase, _RandomForestMixin):

python/cuml/cuml/common/classification.py

@@ -1,5 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
+from collections.abc import Mapping
+
 import cudf
 import cupy as cp
 import numpy as np
@@ -111,6 +113,20 @@ def decode_labels(y_encoded, classes, output_type="cupy", index=None):
         )
 
 
+def validate_class_weight(class_weight):
+    if class_weight is None:
+        return
+    if isinstance(class_weight, str) and class_weight == "balanced":
+        return
+    if isinstance(class_weight, Mapping):
+        return
+
+    raise ValueError(
+        "class_weight must be a dict, 'balanced', or None; "
+        f"got {class_weight!r}"
+    )
+
+
 def process_class_weight(
     classes,
     y_ind,
@@ -152,6 +168,8 @@ def process_class_weight(
     sample_weight: cp.ndarray or None
         The resulting sample weights, or None if uniformly weighted.
     """
+    validate_class_weight(class_weight)
+
     n_classes = len(classes)
     if dtype is None:
         dtype = getattr(sample_weight, "dtype", np.float32)

python/cuml/cuml/dask/ensemble/randomforestclassifier.py

@@ -155,7 +155,13 @@ def _construct_rf(n_estimators, random_state, **kwargs):
             n_estimators=n_estimators, random_state=random_state, **kwargs
         )
 
-    def fit(self, X, y, convert_dtype="deprecated", broadcast_data=False):
+    def fit(
+        self,
+        X,
+        y,
+        convert_dtype="deprecated",
+        broadcast_data=False,
+    ):
         """
         Fit the input data with a Random Forest classifier
 
@@ -207,7 +213,6 @@ def fit(self, X, y, convert_dtype="deprecated", broadcast_data=False):
             When set to True, the whole dataset is broadcasted
             to train the workers, otherwise each worker
             is trained on its partition
-
         """
         # Handle both Dask Arrays and Dask Series/DataFrames
         if isinstance(y, dask.array.Array):

python/cuml/cuml/dask/ensemble/randomforestregressor.py

@@ -138,7 +138,13 @@ def _construct_rf(n_estimators, random_state, **kwargs):
             n_estimators=n_estimators, random_state=random_state, **kwargs
         )
 
-    def fit(self, X, y, convert_dtype="deprecated", broadcast_data=False):
+    def fit(
+        self,
+        X,
+        y,
+        convert_dtype="deprecated",
+        broadcast_data=False,
+    ):
         """
         Fit the input data with a Random Forest regression model
 
@@ -186,7 +192,6 @@ def fit(self, X, y, convert_dtype="deprecated", broadcast_data=False):
             When set to True, the whole dataset is broadcasted
             to train the workers, otherwise each worker
             is trained on its partition
-
         """
         self.internal_model = None
         self._fit(

python/cuml/cuml/ensemble/randomforest_common.pyx

@@ -80,7 +80,8 @@ cdef extern from "cuml/ensemble/randomforest.hpp" namespace "ML" nogil:
         RF_params params,
         bool* bootstrap_masks,
         T* feature_importances,
-        level_enum verbosity
+        level_enum verbosity,
+        const double* sample_weight
     ) except +
 
     cdef void fit_treelite[T, L](
@@ -93,7 +94,8 @@ cdef extern from "cuml/ensemble/randomforest.hpp" namespace "ML" nogil:
         RF_params params,
         bool* bootstrap_masks,
         T* feature_importances,
-        level_enum verbosity
+        level_enum verbosity,
+        const double* sample_weight
     ) except +
 
 
@@ -455,12 +457,15 @@ class BaseRandomForestModel(InteropMixin, Base):
             handle=get_handle(),
         )
 
-    def _fit_forest(self, X, y):
+    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
 
         cdef uintptr_t X_ptr = X.data.ptr
         cdef uintptr_t y_ptr = y.data.ptr
+        cdef uintptr_t sample_weight_ptr = (
+            0 if sample_weight is None else sample_weight.data.ptr
+        )
         cdef int n_rows = X.shape[0]
         cdef int n_cols = X.shape[1]
         cdef level_enum verbose = <level_enum> self._verbose_level
@@ -578,7 +583,8 @@ class BaseRandomForestModel(InteropMixin, Base):
                         params,
                         bootstrap_masks_ptr,
                         <float*> feature_importances_ptr,
-                        verbose
+                        verbose,
+                        <const double*> sample_weight_ptr
                     )
                 else:
                     fit_treelite(
@@ -592,7 +598,8 @@ class BaseRandomForestModel(InteropMixin, Base):
                         params,
                         bootstrap_masks_ptr,
                         <double*> feature_importances_ptr,
-                        verbose
+                        verbose,
+                        <const double*> sample_weight_ptr
                     )
             else:
                 if is_float32:
@@ -606,7 +613,8 @@ class BaseRandomForestModel(InteropMixin, Base):
                         params,
                         bootstrap_masks_ptr,
                         <float*> feature_importances_ptr,
-                        verbose
+                        verbose,
+                        <const double*> sample_weight_ptr
                     )
                 else:
                     fit_treelite(
@@ -619,7 +627,8 @@ class BaseRandomForestModel(InteropMixin, Base):
                         params,
                         bootstrap_masks_ptr,
                         <double*> feature_importances_ptr,
-                        verbose
+                        verbose,
+                        <const double*> sample_weight_ptr
                     )
 
         # XXX: Theoretically we could wrap `tl_handle` with `treelite.Model` to