[MT] Support distributed training and the Dask interface.#12292
Conversation
There was a problem hiding this comment.
Pull request overview
Adds distributed-training support for multi-target (“vector leaf”) hist training (including Dask) by introducing the needed collective reductions for node/leaf statistics and expanding test coverage for both CPU and GPU Dask workflows.
Changes:
- Add new Dask test helpers and new CPU/GPU distributed Dask tests for
multi_output_treewith a value-gradient objective (reg:absoluteerror). - Enable distributed multi-target hist training by globally reducing multi-target root/leaf sums and ensuring histograms are allreduced where required.
- Introduce a dedicated Hypothesis strategy (
hist_multi_parameter_strategy) usage in distributed Dask test suites.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_distributed/test_with_dask/test_with_dask.py | Adds CPU Dask multi-target hist property test + targeted absolute-error and sklearn-style smoke tests. |
| tests/test_distributed/test_gpu_with_dask/test_gpu_with_dask.py | Adds GPU Dask multi-target hist property test + targeted absolute-error and sklearn-style smoke tests (guarded by CuPy availability). |
| src/tree/updater_quantile_hist.cc | Adds collective reduction for multi-target leaf sums (and extra sanity checks) to support distributed vector-leaf training. |
| src/tree/updater_gpu_hist.cuh | Adds collective reductions and required histogram allreduces for multi-target GPU hist, enabling distributed vector-leaf training. |
| python-package/xgboost/testing/dask.py | Adds shared Dask helpers to generate multi-output regression data and validate multi_output_tree training/prediction/config. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
RAMitchell
left a comment
RAMitchell
left a comment
There was a problem hiding this comment.
Just curious, can scalar leaf training be implemented as a base case of vector leaf training? I don't like having two builder implementations.
| if (tree.IsMultiTarget()) { | ||
| tree.SetLeaves(h_node_idx, common::Span{quantiles}); | ||
| } else { | ||
| for (size_t i = 0; i < nidx.size(); ++i) { |
Yes. For now, I would like to keep them separate. We have different evaluators, different tree model formats, etc. I omitted some optimizations for the vector leaf to keep things simple and would like to use it as a base for the CV implementation without breaking or regressing the scalar leaf. |
|
I think life will be easier if you finish this before implementing CV. |
No, then I will have to run benchmarks to avoid regression for every change. |
|
My current WIP CV implementation does use the multi-target code base uniformly with single-target as a trivial case. But unifying the code with the existing scalar leaf would be a huge refactoring effort that I would like to defer until everything is inplace. |
Don't you need to do this eventually? You are currently building a large disruptive feature on top of another disruptive unfinished feature that might still need to change. Is it the most efficient way to work? I won't stop you but just something to think about. |
Yes, but with a clear picture of everything, and can omit benchmarks that don't necessarily matter for the final state.
Thank you for this. Yes, I care deeply about the state of the code. |
ref: #9043