Add distributed RF C++ tree builder core - #8255
Closed
RAMitchell wants to merge 6 commits into
Closed
Conversation
There was a problem hiding this comment.
Pull request overview
This PR adds the C++ core needed for distributed (multi-rank) RandomForest tree building in the batched-level decision tree trainer by all-reducing per-rank histograms to select globally consistent splits, while keeping rank-local row partitioning and preserving the single-GPU behavior when distributed comms aren’t initialized.
Changes:
- Implement distributed histogram aggregation (all-reduce) for split finding and leaf histogram aggregation; track global vs rank-local counts for correct metadata vs local partitioning.
- Update RF training to support ranks with zero local rows, and serialize distributed tree building to ensure consistent collective ordering.
- Add Multi-GPU RF property tests and extend quantile MG test harness; improve feature-importance normalization robustness with infinite split metrics.
Reviewed changes
Copilot reviewed 13 out of 13 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| cpp/tests/sg/rf_test.cu | Adds SG coverage to ensure empty global row training is rejected. |
| cpp/tests/mg/rf_test.cu | New MG RF property tests for distributed consistency across ranks/partitions. |
| cpp/tests/mg/rf_quantile_test.cu | Adds an MPI-aware main() to run MG quantile tests under MPI. |
| cpp/tests/CMakeLists.txt | Registers the new MG RF test target. |
| cpp/src/randomforest/randomforest.cuh | Enables distributed-mode behavior (allow empty local rows, force serial trees, ensure correct device in OMP threads). |
| cpp/src/randomforest/randomforest.cu | Makes feature-importance normalization robust to infinite split contributions. |
| cpp/src/decisiontree/batched-levelalgo/split.cuh | Extends split metadata to include global + local left counts using 64-bit counters. |
| cpp/src/decisiontree/batched-levelalgo/quantiles.cuh | Allows n_rows==0 locally in distributed mode; guards null data pointer accordingly. |
| cpp/src/decisiontree/batched-levelalgo/objectives.cuh | Promotes split-count math to 64-bit and adapts objectives to updated split representation. |
| cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh | Introduces split-histogram/evaluation kernel split and histogram pack/unpack helpers for all-reduce. |
| cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels_impl.cuh | Implements the two-phase histogram build + split evaluation flow; updates node split/leaf logic for global+local counts. |
| cpp/src/decisiontree/batched-levelalgo/builder.cuh | Adds distributed all-reduce for histograms and leaf histograms; wires global sampled-row counts into tree metadata. |
| cpp/src/decisiontree/batched-levelalgo/bins.cuh | Switches histogram bin counts to 64-bit integer counters to avoid overflow. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+168
to
+183
| for (size_t i = 0; i < rows.size(); ++i) { | ||
| int global_row = rows[i]; | ||
| DataT signal = static_cast<DataT>((global_row % 97) - 48); | ||
| for (int col = 0; col < params.n_cols; ++col) { | ||
| DataT feature = signal * static_cast<DataT>(col + 1); | ||
| feature += static_cast<DataT>(((global_row + 13 * col + params.seed) % 11) - 5) / | ||
| static_cast<DataT>(10); | ||
| X[i * params.n_cols + col] = feature; | ||
| } | ||
| if constexpr (std::is_integral_v<LabelT>) { | ||
| y[i] = (signal >= DataT(0)) ? 1 : 0; | ||
| if (params.n_labels > 2 && global_row % 17 == 0) { y[i] = 2; } | ||
| } else { | ||
| y[i] = signal * DataT(0.5) + static_cast<DataT>((global_row % 7) - 3); | ||
| } | ||
| } |
Contributor
Author
There was a problem hiding this comment.
Addressed in 3f06f4e: make_local_dataset now writes the training buffer in column-major layout with X[col * rows.size() + i], matching fit() / the batched-level tree builder. The prediction path remains row-major and no longer transposes before predict().
Comment on lines
+327
to
+333
| rmm::device_uvector<DataT> X(h_X.size(), handle.get_stream()); | ||
| rmm::device_uvector<DataT> X_transpose(h_X.size(), handle.get_stream()); | ||
| rmm::device_uvector<LabelT> predictions(params.n_rows, handle.get_stream()); | ||
| raft::update_device(X.data(), h_X.data(), h_X.size(), handle.get_stream()); | ||
| raft::linalg::transpose( | ||
| handle, X.data(), X_transpose.data(), params.n_rows, params.n_cols, handle.get_stream()); | ||
| predict(handle, forest, X_transpose.data(), params.n_rows, params.n_cols, predictions.data()); |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughSummary by CodeRabbitRelease Notes
WalkthroughRefactors batched decision-tree and RandomForest training for distributed multi-GPU execution. ChangesDistributed Random Forest Training
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
cpp/src/decisiontree/batched-levelalgo/split.cuh (1)
1-4:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winUpdate copyright year to 2026.
The copyright header shows
2019-2022but should reflect the current year2019-2026.Proposed fix
/* - * SPDX-FileCopyrightText: Copyright (c) 2019-2022, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */🤖 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/split.cuh` around lines 1 - 4, Update the file header comment block at the top of split.cuh by changing the copyright year range from "2019-2022" to "2019-2026" so the SPDX header reads "SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION." and keep the SPDX-License-Identifier unchanged.Source: Coding guidelines
🧹 Nitpick comments (2)
cpp/src/decisiontree/batched-levelalgo/objectives.cuh (1)
283-288: 💤 Low valueConsider adding division-by-zero guard in
SetLeafVector.If
shist[i].countis zero, this produces Inf/NaN. While leaf nodes should always have samples (upstream logic should prevent this), a defensive epsilon or guard would improve robustness.Optional defensive fix
static DI void SetLeafVector(BinT const* shist, int nclasses, DataT* out) { for (int i = 0; i < nclasses; i++) { - out[i] = shist[i].label_sum / shist[i].count; + out[i] = shist[i].count > 0 ? shist[i].label_sum / shist[i].count : DataT(0); } }🤖 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 283 - 288, SetLeafVector can produce Inf/NaN when shist[i].count == 0; update the loop in SetLeafVector to guard the division by checking shist[i].count for zero and handle it defensively (e.g., if count == 0 set out[i] to 0 (or a safe default) otherwise compute out[i] = shist[i].label_sum / shist[i].count); reference the SetLeafVector function and the shist[i].count/shist[i].label_sum/out arrays when making this change (alternatively you can use a small epsilon divisor instead of the zero-check if preferred).cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh (1)
402-405: Remove unusedhistogramTransformBlockshelper
histogramTransformBlocksis only defined incpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh(around lines 402-405) and has no other references/call sites in the repository; remove it or inline the expression where needed.🤖 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/kernels/builder_kernels.cuh` around lines 402 - 405, Remove the unused helper histogramTransformBlocks: delete its definition and, if you intended to use it elsewhere, replace any future call sites with the equivalent expression std::max<std::size_t>(std::size_t{1}, raft::ceildiv<std::size_t>(len, 256)); otherwise simply remove the histogramTransformBlocks function declaration/definition to eliminate dead code (search for histogramTransformBlocks to ensure no references remain).
🤖 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`:
- Around line 244-252: The code incorrectly uses the global MPI rank and global
size for CUDA device selection: compute a node-local rank/size and use that to
validate and set the device instead of using the global rank/size. Replace the
comparison of n_gpus against size with a check against the per-node local_size,
and call cudaSetDevice(local_rank) (wrapped with RAFT_CUDA_TRY) where local_rank
is the process's rank on the local node; obtain local_rank/local_size by
creating a node-local communicator (e.g., MPI_Comm_split_type with
MPI_COMM_TYPE_SHARED) or by reading standard local-rank environment variables,
so that n_gpus, cudaGetDeviceCount, cudaSetDevice and RAFT_CUDA_TRY use
node-local indices rather than the global rank/size.
---
Outside diff comments:
In `@cpp/src/decisiontree/batched-levelalgo/split.cuh`:
- Around line 1-4: Update the file header comment block at the top of split.cuh
by changing the copyright year range from "2019-2022" to "2019-2026" so the SPDX
header reads "SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA
CORPORATION." and keep the SPDX-License-Identifier unchanged.
---
Nitpick comments:
In `@cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh`:
- Around line 402-405: Remove the unused helper histogramTransformBlocks: delete
its definition and, if you intended to use it elsewhere, replace any future call
sites with the equivalent expression std::max<std::size_t>(std::size_t{1},
raft::ceildiv<std::size_t>(len, 256)); otherwise simply remove the
histogramTransformBlocks function declaration/definition to eliminate dead code
(search for histogramTransformBlocks to ensure no references remain).
In `@cpp/src/decisiontree/batched-levelalgo/objectives.cuh`:
- Around line 283-288: SetLeafVector can produce Inf/NaN when shist[i].count ==
0; update the loop in SetLeafVector to guard the division by checking
shist[i].count for zero and handle it defensively (e.g., if count == 0 set
out[i] to 0 (or a safe default) otherwise compute out[i] = shist[i].label_sum /
shist[i].count); reference the SetLeafVector function and the
shist[i].count/shist[i].label_sum/out arrays when making this change
(alternatively you can use a small epsilon divisor instead of the zero-check if
preferred).
🪄 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: bfe64907-3651-42e0-a9cf-4350bd65ae81
📒 Files selected for processing (13)
cpp/src/decisiontree/batched-levelalgo/bins.cuhcpp/src/decisiontree/batched-levelalgo/builder.cuhcpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuhcpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels_impl.cuhcpp/src/decisiontree/batched-levelalgo/objectives.cuhcpp/src/decisiontree/batched-levelalgo/quantiles.cuhcpp/src/decisiontree/batched-levelalgo/split.cuhcpp/src/randomforest/randomforest.cucpp/src/randomforest/randomforest.cuhcpp/tests/CMakeLists.txtcpp/tests/mg/rf_quantile_test.cucpp/tests/mg/rf_test.cucpp/tests/sg/rf_test.cu
RAMitchell
added
improvement
…f-cpp-core # Conflicts: # cpp/src/decisiontree/batched-levelalgo/bins.cuh # 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/objectives.cuh
…f-cpp-core # Conflicts: # 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
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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/kernels/builder_kernels_impl.cuh`:
- Around line 78-81: `nodeSplitLocalCountKernel` and the downstream index
calculation in `builder_kernels_impl.cuh` are using `split.local_nLeft` as if it
were always zero, but the count is only accumulated via `atomicAdd` into
`splits[nid].local_nLeft`. Reset `local_nLeft` before running the local count
kernel, or ensure the selected split copied into the build path is initialized
with `local_nLeft == 0`, so `rank`, `local_left_count`, and `out_idx` stay
correct for right-side rows. Apply the same fix consistently in the related code
paths referenced by `nodeSplitLocalCountKernel` and the split-building logic.
In `@cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh`:
- Around line 60-64: The split validation in SplitNotValid currently only checks
colid and a direct best_metric_val comparison, so NaN gains can still pass as
valid. Update SplitNotValid in builder_kernels.cuh and the matching builder path
check in builder.cuh to reject non-finite split gains (or use a NaN-safe
comparison) before a split is accepted or written into the tree. Keep the fix
centered on SplitNotValid and the builder validation logic so both paths behave
consistently.
🪄 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: 48c48b8d-ad11-447f-95b4-c1336b5a2bb3
📒 Files selected for processing (5)
cpp/src/decisiontree/batched-levelalgo/builder.cuhcpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuhcpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels_impl.cuhcpp/src/randomforest/randomforest.cucpp/tests/sg/rf_test.cu
🚧 Files skipped from review as they are similar to previous changes (2)
- cpp/src/randomforest/randomforest.cu
- cpp/src/decisiontree/batched-levelalgo/builder.cuh
Comment on lines
+78
to
+81
| const auto rank = state.goes_left ? std::size_t(state.left_count - CountT{1}) | ||
| : range_pos - std::size_t(state.left_count); | ||
| const auto local_left_count = std::size_t(split.local_nLeft); | ||
| const auto out_idx = range_start + (state.goes_left ? rank : local_left_count + rank); |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
HIGH: Reset local_nLeft before accumulating local counts.
nodeSplitLocalCountKernel only atomicAdds into splits[nid].local_nLeft; if the selected global split carries a non-zero or stale local count, Lines 78-81 offset right-side rows incorrectly. Add a reset launch before the count kernel, or verify the copied split always has local_nLeft == 0.
Suggested fix
+template <typename DataT>
+static __global__ void resetNodeSplitLocalCountKernel(Split<DataT>* splits,
+ const WorkloadInfo* workload_info)
+{
+ using CountT = typename Split<DataT>::CountT;
+ const auto nid = workload_info[blockIdx.x].nodeid;
+ splits[nid].local_nLeft = CountT{0};
+}
+
template <typename DataT, typename LabelT, typename IdxT, int TPB>
static __global__ void nodeSplitLocalCountKernel(const DataT min_impurity_decrease,
const Dataset<DataT, LabelT, IdxT> dataset,
@@
using CountT = typename Split<DataT>::CountT;
+ resetNodeSplitLocalCountKernel<DataT><<<n_blocks_dimx, 1, 0, builder_stream>>>(
+ splits, workload_info);
nodeSplitLocalCountKernel<DataT, LabelT, IdxT, TPB><<<n_blocks_dimx, TPB, 0, builder_stream>>>(
min_impurity_decrease, dataset, work_items, splits, workload_info);Also applies to: 87-108, 147-149
🤖 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/kernels/builder_kernels_impl.cuh`
around lines 78 - 81, `nodeSplitLocalCountKernel` and the downstream index
calculation in `builder_kernels_impl.cuh` are using `split.local_nLeft` as if it
were always zero, but the count is only accumulated via `atomicAdd` into
`splits[nid].local_nLeft`. Reset `local_nLeft` before running the local count
kernel, or ensure the selected split copied into the build path is initialized
with `local_nLeft == 0`, so `rank`, `local_left_count`, and `out_idx` stay
correct for right-side rows. Apply the same fix consistently in the related code
paths referenced by `nodeSplitLocalCountKernel` and the split-building logic.
Comment on lines
+60
to
64
| template <typename SplitT, typename DataT> | ||
| HDI bool SplitNotValid(const SplitT& split, DataT min_impurity_decrease) | ||
| { | ||
| return split.best_metric_val <= min_impurity_decrease || | ||
| SplitPartitionNotValid(split, min_samples_leaf, num_rows); | ||
| return split.colid == -1 || split.best_metric_val <= min_impurity_decrease; | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== builder_kernels.cuh (around SplitNotValid) ==\n'
nl -ba cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh | sed -n '1,180p'
printf '\n== usages of SplitNotValid ==\n'
rg -n "SplitNotValid|best_metric_val" cpp/src/decisiontree -g '!**/build/**'Repository: rapidsai/cuml
Length of output: 241
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== builder_kernels.cuh (around SplitNotValid) ==\n'
nl -ba cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh | sed -n '1,180p'
printf '\n== usages of SplitNotValid ==\n'
rg -n "SplitNotValid|best_metric_val" cpp/src/decisiontree -g '!**/build/**'Repository: rapidsai/cuml
Length of output: 241
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== builder_kernels.cuh (around SplitNotValid) ==\n'
sed -n '1,180p' cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh | cat -n
printf '\n== usages of SplitNotValid ==\n'
rg -n "SplitNotValid|best_metric_val" cpp/src/decisiontree -g '!**/build/**'Repository: rapidsai/cuml
Length of output: 10841
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== builder.cuh around line 100 ==\n'
sed -n '80,130p' cpp/src/decisiontree/batched-levelalgo/builder.cuh | cat -n
printf '\n== split.cuh around metric initialization/comparison ==\n'
sed -n '1,220p' cpp/src/decisiontree/batched-levelalgo/split.cuh | cat -nRepository: rapidsai/cuml
Length of output: 10767
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== builder.cuh around line 100 ==\n'
sed -n '80,130p' cpp/src/decisiontree/batched-levelalgo/builder.cuh | cat -n
printf '\n== split.cuh around metric initialization/comparison ==\n'
sed -n '1,220p' cpp/src/decisiontree/batched-levelalgo/split.cuh | cat -nRepository: rapidsai/cuml
Length of output: 10767
Handle NaN split gains in the builder path
cpp/src/decisiontree/batched-levelalgo/kernels/builder_kernels.cuh:61-63 treats NaN as valid, and cpp/src/decisiontree/batched-levelalgo/builder.cuh:21 has the same direct compare. That lets invalid gains slip through and be written into the tree; reject non-finite gains here or use a NaN-safe comparison in both places.
🤖 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/kernels/builder_kernels.cuh` around
lines 60 - 64, The split validation in SplitNotValid currently only checks colid
and a direct best_metric_val comparison, so NaN gains can still pass as valid.
Update SplitNotValid in builder_kernels.cuh and the matching builder path check
in builder.cuh to reject non-finite split gains (or use a NaN-safe comparison)
before a split is accepted or written into the tree. Keep the fix centered on
SplitNotValid and the builder validation logic so both paths behave
consistently.
Source: Coding guidelines
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
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.
Summary
This PR implements the next C++ stage of distributed Random Forest training: workers now collaborate on split finding by all-reducing per-rank histograms, then applying the same globally selected split while partitioning only local row IDs.
This builds on the distributed quantile work in #8111 and #8190. It is a step toward #7969, and is related to the existing multi-GPU RF accuracy reports in #4429, #2437, and #4740.
Refs #7969
Refs #4429
Refs #2437
Refs #4740
Changes
Notes
This PR does not yet add the Python/Dask estimator interface wiring. That should be a follow-up stage after the C++ distributed builder path lands.
Testing
git diff --checkcmake --build cpp/build --target MG_RF_TEST MG_RF_QUANTILE_TEST -j8mpiexec -np 2 ./cpp/build/tests/MG_RF_QUANTILE_TESTmpiexec -np 2 ./cpp/build/tests/MG_RF_TESTmpiexec -np 4 ./cpp/build/tests/MG_RF_QUANTILE_TESTmpiexec -np 4 ./cpp/build/tests/MG_RF_TEST