Distributed RF: add gtests, handle empty partitions, fix bug in leaf output

[RAMitchell]

Description

Adds C++ multi-GPU coverage for the distributed random forest training path and fixes the local/global count bookkeeping needed for ranks with uneven or empty local partitions.

This is preparation for enabling distributed random forest end to end. The tests exercise distributed histogram/split selection with balanced, imbalanced, and empty-rank row partitions, while keeping local partition ranges rank-local.

Key changes:

• Add an MG_RF_TEST target.
• Add distributed RF property tests for classification and regression.
• Allow distributed ranks with zero local rows to participate in quantile computation.
• Track global sampled row counts for tree node counts and expansion decisions.
• Keep local_nLeft scoped to local partition range updates.
• Preserve single-GPU validation that zero training rows are rejected.
• Add missing allreduce before leaf output computation.
• Correctly handle sample weights when some partitions are empty.
• Compare the RF model objects from the multi-GPU and single-GPU algorithms, using a suitable tolerances for floating-point values.

Validation

cmake --build /home/rorym/cuml-builds/codex-enh-rf-mg-tests/cpp-mg-test-2610-mpi --target MG_RF_TEST -j8

OMPI_MCA_opal_cuda_support=true UCX_MEMTYPE_CACHE=n NCCL_P2P_DISABLE=1 NCCL_SHM_DISABLE=1 \
  mpiexec -np 2 /home/rorym/cuml-builds/codex-enh-rf-mg-tests/cpp-mg-test-2610-mpi/tests/MG_RF_TEST

Result: 8/8 tests passed.

/home/rorym/cuml-builds/codex-enh-rf-mg-tests/cpp-mg-test-2610-mpi/tests/SG_RF_TEST \
  --gtest_filter='RfTests.EmptyGlobalRowsRejected'

Result: passed.

pre-commit run --files \
  cpp/src/decisiontree/batched-levelalgo/builder.cuh \
  cpp/src/decisiontree/batched-levelalgo/quantiles.cuh \
  cpp/src/randomforest/randomforest.cuh \
  cpp/tests/CMakeLists.txt \
  cpp/tests/sg/rf_test.cu \
  cpp/tests/mg/rf_test.cu

Result: passed.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[chyunsu3]

Counterexample: a single tree stump

n_rows=128
n_cols=4
n_trees=1
max_features=1.0f
max_depth=1
max_leaves=-1
max_n_bins=16
min_samples_leaf=1
min_samples_split=2
min_impurity_decrease=0.0f
n_streams=1
handle_n_streams=1
split_criterion=GINI
seed=7
n_labels=2
double_precision=false
partition_kind=PartitionKind::Contiguous

Single-GPU:

 Decision Tree depth --> 1 and n_leaves --> 2
 Tree Fitting - Overall time --> 4.437 milliseconds
└(colid: 3, quesval: -4, best_metric_val: 0.472534)
    ├(leaf, prediction: [1, 0], best_metric_val: 0)
    └(leaf, prediction: [0, 1], best_metric_val: 0)

4 GPUs: the workers disagree on the content of leaf output

 Decision Tree depth --> 1 and n_leaves --> 2
 Tree Fitting - Overall time --> 8.481 milliseconds
└(colid: 3, quesval: -4, best_metric_val: 0.472534)
    ├(leaf, prediction: [1, 0], best_metric_val: 0)
    └(leaf, prediction: [0, 0], best_metric_val: 0)

 Decision Tree depth --> 1 and n_leaves --> 2
 Tree Fitting - Overall time --> 8.736 milliseconds
└(colid: 3, quesval: -4, best_metric_val: 0.472534)
    ├(leaf, prediction: [0, 0], best_metric_val: 0)
    └(leaf, prediction: [0, 1], best_metric_val: 0)

 Decision Tree depth --> 1 and n_leaves --> 2
 Tree Fitting - Overall time --> 7.651 milliseconds
└(colid: 3, quesval: -4, best_metric_val: 0.472534)
    ├(leaf, prediction: [1, 0], best_metric_val: 0)
    └(leaf, prediction: [0, 1], best_metric_val: 0)

 Decision Tree depth --> 1 and n_leaves --> 2
 Tree Fitting - Overall time --> 7.95 milliseconds
└(colid: 3, quesval: -4, best_metric_val: 0.472534)
    ├(leaf, prediction: [1, 0], best_metric_val: 0)
    └(leaf, prediction: [0, 1], best_metric_val: 0)

[chyunsu3]

Fix: Run an AllReduce on the statistics histogram before computing the leaf value.

[RAMitchell]

New changes look correct!

[RAMitchell]

One more thing:

Distributed n_rows == 0 is now allowed, but in weighted bootstrap, compute_sample_weight_sum() reads sample_weight_cdf_.data() + n_rows_ - 1. This is out of bounds.

For non-bootstrap the reduction returns 0.0 and then fails local sample_weight_sum_ > 0.0.

So it looks like we need coverage for a worker with 0 rows for weighted and unweighted.

[chyunsu3]

TODOs

• Inspect coderabbtai reviews and see if they are valid.
• Improve test coverage with empty partition

[csadorf]

Looks mostly good to me, just some minor concerns.

[chyunsu3]

I addressed all review comments. Can you take another look? @RAMitchell @csadorf

[coderabbitai]

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

cpp/src/randomforest/randomforest.cuh (1)

120-120: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Early return should check n_rows_ == 0 instead of selected_rows.size() == 0.

The validation at cpp/src/randomforest/randomforest.cu:548–550 confirms that max_samples is permitted in the range (0, 1]. For small partitions with small max_samples values, round(max_samples * n_rows) produces zero. For example, round(0.1 * 5) = 0. In such cases, n_rows_ is non-zero but n_sampled_rows_ is zero, so selected_rows.size() == 0 is true even though the partition is not empty.

The early return at line 120 skips the subsequent call to store_bootstrap_mask() at line 168. If bootstrap_masks_ is non-null, the bootstrap mask for that tree is never initialized, leaving the caller's buffer uninitialized or stale.

The store_bootstrap_mask() method safely handles null pointers at line 180, so it is safe to call unconditionally. Change the condition to if (n_rows_ == 0) { return selected_rows; } to skip initialization only for empty partitions.

🤖 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/randomforest/randomforest.cuh` at line 120, Update the early-return
condition in the random-forest sampling flow to check n_rows_ == 0 instead of
selected_rows.size() == 0. Preserve execution of store_bootstrap_mask() when
n_rows_ is nonzero but sampling produces zero rows, while retaining the return
for genuinely empty partitions.

Source: Path instructions

🧹 Nitpick comments (2)

cpp/tests/mg/rf_test.cu (2)

291-301: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Consider zero-filling the placeholder weight element.

Line 296 allocates one double for an empty rank to keep the pointer non-null. Line 299 copies only h_sample_weights.size() elements, so that element keeps uninitialized device memory. The builder should not read it when the local row count is zero, but a zero-filled buffer removes the dependency on that assumption and makes a future regression fail deterministically instead of nondeterministically.

♻️ Proposed change

     rmm::device_uvector<double> sample_weights(sample_weight_buffer_size, handle.get_stream());
+    RAFT_CUDA_TRY(cudaMemsetAsync(sample_weights.data(),
+                                  0,
+                                  sample_weights.size() * sizeof(double),
+                                  handle.get_stream()));
     raft::update_device(X.data(), h_X.data(), h_X.size(), handle.get_stream());

🤖 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/tests/mg/rf_test.cu` around lines 291 - 301, Zero-initialize the
placeholder element in sample_weights when params.use_sample_weights is true and
h_sample_weights is empty, while preserving the existing host-to-device copy for
actual weights. Update the sample_weights allocation or initialization near
sample_weight_buffer_size and keep sample_weight_ptr behavior unchanged.

---

640-657: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add a weighted classification case.

use_sample_weights is true only for the final MSE case. The stack also changes the weighted classification kernels, and the reported four-GPU leaf-value counterexample was a classification stump. One weighted GINI entry with PartitionKind::EmptyNonRootRanks would cover that path with the same fixture. Apply the global_row-keyed weight fix first, so a weighted Strided or Imbalanced entry also stays valid.

🤖 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/tests/mg/rf_test.cu` around lines 640 - 657, Add a weighted
classification test case alongside the existing parameterized cases, using GINI
with use_sample_weights enabled and PartitionKind::EmptyNonRootRanks while
preserving the fixture’s expected values. Ensure the global_row-keyed weight fix
is applied so the new weighted Strided or Imbalanced variants remain valid.

🤖 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 `@cpp/tests/mg/rf_test.cu`:
- Line 197: Update the sample-weight assignment in make_local_dataset to derive
the parity from global_row rather than the local index i, so distributed and
reconstructed single-node datasets assign identical weights regardless of rank
partition sizes.

---

Outside diff comments:
In `@cpp/src/randomforest/randomforest.cuh`:
- Line 120: Update the early-return condition in the random-forest sampling flow
to check n_rows_ == 0 instead of selected_rows.size() == 0. Preserve execution
of store_bootstrap_mask() when n_rows_ is nonzero but sampling produces zero
rows, while retaining the return for genuinely empty partitions.

---

Nitpick comments:
In `@cpp/tests/mg/rf_test.cu`:
- Around line 291-301: Zero-initialize the placeholder element in sample_weights
when params.use_sample_weights is true and h_sample_weights is empty, while
preserving the existing host-to-device copy for actual weights. Update the
sample_weights allocation or initialization near sample_weight_buffer_size and
keep sample_weight_ptr behavior unchanged.
- Around line 640-657: Add a weighted classification test case alongside the
existing parameterized cases, using GINI with use_sample_weights enabled and
PartitionKind::EmptyNonRootRanks while preserving the fixture’s expected values.
Ensure the global_row-keyed weight fix is applied so the new weighted Strided or
Imbalanced variants remain valid.

🪄 Autofix

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: 10dde977-a3b6-48b2-9526-8cfff82c5896

📥 Commits

Reviewing files that changed from the base of the PR and between db8418b and efa28f1.

📒 Files selected for processing (2)

• cpp/src/randomforest/randomforest.cuh
• cpp/tests/mg/rf_test.cu

[RAMitchell]

A few more from codex:

Findings
[P2] randomforest.cuh (line 120): RowSampler::sample() now returns on selected_rows.size() == 0. That fixes empty local ranks, but it also catches non-empty ranks where max_samples * n_rows rounds to zero. In that case store_bootstrap_mask() is skipped, so an OOB/bootstrap mask can remain uninitialized or stale. This should be if (n_rows_ == 0).
[P2] rf_test.cu (line 274): the GPU-count failure still returns rank-locally before later collective paths. If only some ranks hit this, the rest can hang. This was one of the review concerns and is still present.
[P2] rf_test.cu (line 197): sample weights are still keyed by local i, not global_row. Current weighted empty-rank case happens to match because rank 0 owns all rows, but the helper is fragile for weighted strided/imbalanced cases.
[P3] rf_test.cu (line 220): child-count checks still use LeftChildId() + 1 and no bounds checks. Better to use RightChildId() and assert both IDs before indexing.
[P3] CMakeLists.txt (line 232): MG_RF_TEST is registered with the default GPUS 1 and no MPI launcher in ConfigureTest. If CI wraps this externally, fine, but the specific scheduling concern is not addressed in the PR itself.
Addressed
The original empty-rank weighted sample-weight concern is mostly fixed: randomforest.cuh (line 94) skips local weight-sum validation for n_rows_ == 0, and the test forces a non-null weight pointer for empty weighted ranks.
Regression forest comparison was improved to use tolerances rather than raw bitwise hashing.
C++ build/style/test jobs on CI are green so far; many Python/wheel jobs are still pending. GitHub still shows CHANGES_REQUESTED, matching the unresolved threads above.

[csadorf]

Please address @RAMitchell 's and the CodeRabbit comments.

[csadorf]

@chyunsu3 can you also update the PR title? I don't think it's fully matching the PR scope and intent anymore.

[chyunsu3]

I addressed all review comments and updated the title and the description.

[RAMitchell]

Thanks, LGTM

[chyunsu3]

/merge

[csadorf]

/merge