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!

[coderabbitai]

Actionable comments posted: 7

🧹 Nitpick comments (2)

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

149-158: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Document the column-major layout of the generated X.

Line 158 writes X in column-major order with the leading dimension set to the local row count. The layout is implied only by the index expression. State it in a comment so the helper does not depend on an implicit assumption, and so the match with the fit default layout stays visible.

♻️ Proposed comment

 {
+  // X is column-major with leading dimension equal to the local row count,
+  // matching the default (non row-major) layout expected by fit().
   X.resize(rows.size() * params.n_cols);
   y.resize(rows.size());

As per coding guidelines: "Function parameters with ambiguous data format (row-major or column-major) must be explicitly documented or validated at the function entry point".

🤖 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 149 - 158, Add an explicit comment at
the start of the generated X population in the relevant test helper, documenting
that X uses column-major storage with the local row count as its leading
dimension, matching fit’s default layout. Keep the existing indexing and data
generation unchanged.

Source: Coding guidelines

---

356-360: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Report the diverging rank index in the failure message.

The loop compares every hash to hashes.front(). When a mismatch occurs, the message names only the label. Include the rank index so a failure in a multi-rank run points at the diverging rank.

♻️ Proposed refactor

-    for (auto hash : hashes) {
-      EXPECT_EQ(hash, hashes.front()) << "Mismatched distributed RF " << label;
-    }
+    for (std::size_t r = 0; r < hashes.size(); ++r) {
+      EXPECT_EQ(hashes[r], hashes.front())
+        << "Mismatched distributed RF " << label << " on rank " << r;
+    }

🤖 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 356 - 360, Update the hash-comparison
loop around hashes and hashes.front() to track each rank’s index and include it
in the EXPECT_EQ failure message, while preserving the existing label and
comparison behavior.

🤖 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/src/decisiontree/batched-levelalgo/builder.cuh`:
- Around line 684-690: Convert the computed leaf batch count to int exactly once
using ML::narrow_cast<int> before the launcher calls. Update the affected calls
in the leaf-statistics flow around max_batch_size and lines 704–721 to reuse
that checked int value, avoiding implicit narrowing while preserving the
existing batch-size calculation.

In `@cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels_impl.cuh`:
- Around line 245-258: Update both CUDA kernel launchers, including
launchBuildLeafHistogramsKernel and the launcher around the second referenced
block, so every grid or block dimension passed directly to <<<...>>> is wrapped
with ML::narrow_cast<ML::cuda_launch_t>(...). Apply the conversion to num_blocks
and any other direct launch dimensions while preserving the existing launch
configuration.

In `@cpp/src/randomforest/randomforest.cuh`:
- Around line 239-253: Update RowSampler’s sample-weight processing to bypass
CDF construction and compute_sample_weight_sum() when n_rows_ == 0, regardless
of bootstrap mode, avoiding access to the final CDF element and zero-sum
assertions. Keep initialization of the zero-sized selected-row buffers
unchanged, and preserve existing processing for non-empty ranks.

In `@cpp/tests/mg/rf_test.cu`:
- Around line 324-347: Update the distributed-versus-single-node assertion
around hash_forest_structure to avoid bitwise comparison for regression
criteria, whose floating-point reductions are not reproducible and whose
single-node fit changes n_streams. Restrict the hash equality check to
classification criteria, following the existing regression-test handling in
rf_test.cu; alternatively, use structure and threshold comparisons with an
appropriate tolerance for regression models.
- Around line 186-192: Update the non-leaf child-count assertions in the
sparsetree loop to obtain the right child through node.RightChildId() instead of
deriving it from node.LeftChildId() + 1. Before indexing tree->sparsetree,
assert that both child IDs are valid bounds, following the existing single-GPU
test pattern, then preserve the InstanceCount equality assertion.
- Around line 240-247: Update the GPU-count validation before the
`cudaSetDevice` call to communicate the failure condition across all MPI ranks
with an all-reduce, ensuring every rank exits consistently before later
collectives when any rank has insufficient GPUs. Apply the same synchronized
abort decision to the reconstruction check around
`expect_identical_across_ranks`, replacing rank-local early returns while
preserving the existing failure reporting.
- Around line 493-499: Update the MG_RF_TEST configuration associated with main
so it explicitly requests the required multi-GPU allocation and registers the
MPI launcher, guarded by MPI availability. Do not rely on the MPIEnvironment
setup in main to provide GPU scheduling; preserve the existing test execution
flow.

---

Nitpick comments:
In `@cpp/tests/mg/rf_test.cu`:
- Around line 149-158: Add an explicit comment at the start of the generated X
population in the relevant test helper, documenting that X uses column-major
storage with the local row count as its leading dimension, matching fit’s
default layout. Keep the existing indexing and data generation unchanged.
- Around line 356-360: Update the hash-comparison loop around hashes and
hashes.front() to track each rank’s index and include it in the EXPECT_EQ
failure message, while preserving the existing label and comparison behavior.

🪄 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: d6dc6a8f-fb55-4cd0-a548-027ec9972169

📥 Commits

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

📒 Files selected for processing (16)

• cpp/src/decisiontree/batched-levelalgo/builder.cuh
• cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh
• cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels_impl.cuh
• cpp/src/decisiontree/batched-levelalgo/kernels/classification-double.cu
• cpp/src/decisiontree/batched-levelalgo/kernels/classification-float.cu
• cpp/src/decisiontree/batched-levelalgo/kernels/regression-double.cu
• cpp/src/decisiontree/batched-levelalgo/kernels/regression-float.cu
• cpp/src/decisiontree/batched-levelalgo/kernels/weighted-classification-double.cu
• cpp/src/decisiontree/batched-levelalgo/kernels/weighted-classification-float.cu
• cpp/src/decisiontree/batched-levelalgo/kernels/weighted-regression-double.cu
• cpp/src/decisiontree/batched-levelalgo/kernels/weighted-regression-float.cu
• cpp/src/decisiontree/batched-levelalgo/quantiles.cuh
• cpp/src/randomforest/randomforest.cuh
• cpp/tests/CMakeLists.txt
• cpp/tests/mg/rf_test.cu
• cpp/tests/sg/rf_test.cu

[csadorf]

@chyunsu3 you marked this PR as ready for review, but the description still says it's WIP. Is this RP ready? If so, please update the description.

[chyunsu3]

Updated the description.

[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