Ray Distributed XGBoost on GPU
Garrett Frere
XGBoost is Serial and CPU Bound
XGBoost is inherently iterative, the algorithm iterates through boosting rounds making serial updates based on the previous iterations. For binary classification, each boosting round produces a single tree whose structure depends on the residuals of all previous trees — round N+1 cannot begin until round N completes. Within a single round, tree construction is recursive branching: evaluate splits, partition data, recurse into children. This is sequential, narrow, and CPU-appropriate by design.
Let’s explore what “narrow” means. Binary classification: 10 million rows, 50 features, max_bin=256, n_estimators=500. The gradient array is 10M × 2 × 4 bytes = 80 MB — compact enough to reside in CPU cache. The factor of 2 accounts for one gradient (first derivative) and one hessian (second derivative) per sample, each stored as a 4-byte float. Histogram construction evaluates 256 bins × 50 features = 12,800 accumulation targets per split. A few CPU cores complete this comfortably. In distributed Ray form like when using Snowflake’s XGB Distributor, an AllReduce sums each worker's local histograms into a global histogram before the split decision is made — but with payloads this small, orchestration is negligible.
The work is narrow in every dimension: the gradient array is small because there are only two classes, the histogram is small because default bins suffice, and the model converges quickly because binary separation is a simpler learning problem. Early stopping triggers quickly — with only two classes, the ensemble exhausts useful splits sooner than with more potential target classes and evaluation loss plateaus, and training completes quickly. Finally, gradients are overwritten in place with each round, so peak memory doesn’t expand significantly beyond the initial allocation.
XGBoost workloads that have this profile run efficiently on CPU optimized compute. Space complexity is low as gradients are overwritten each round, memory requirements scale from training set size, and the operation is too narrow to benefit from massive parallelism.
What the Inner Loop Actually Does
Each boosting iteration builds a tree. Building a tree means finding the best split at every node. Finding the best split means evaluating every feature at every bin boundary. This is histogram construction — the innermost and most repeated operation in XGBoost training. It executes once per node, per split candidate, per tree, per round.
for each leaf node being split:
for each sample in this node:
for each feature:
look up which bin this sample falls into
accumulate gradient[sample][class] into
histogram[feature][bin][class]The operation is simple — read a gradient value, add it to a histogram bin. There is no complex math. The bottleneck is how fast memory can feed data to the cores, not how fast the cores can process it.
The width of this operation — how many independent accumulations exist per split — determines whether CPU or GPU is the correct hardware. When the number of classes changes, so does the width.
When the Shape Changes
Now consider a multi-class problem. 400+ target classes with a long-tailed distribution — the top 20 classes hold 60% of samples, the bottom 200 have fewer than 500 samples each. This is common in payment classification, product taxonomy, and fraud detection with granular categories.
The gradient array explodes. Each sample now produces a gradient and hessian per class — 800 values instead of 2. For 50M samples that’s 50M × 400 × 2 × 4 bytes = 160 GB. It no longer fits on cheap single node VMs without expensive high-memory provisioning.
The bin resolution must also increase. max_bin controls how many discrete boundaries XGBoost evaluates per feature when searching for splits. The default of 256 is sufficient when classes are well-separated, but with 400+ classes competing for splits in the same feature space, narrow overlap regions between minority classes require finer granularity. Practitioners push to 1000-2000+ bins to get acceptable recall on tail classes — a compensating mechanism for imbalance that also multiplies the parallel work in the inner loop.
Unlike the binary case, where evaluation loss plateaus quickly and early stopping ends training in a few rounds, multi-class with long-tailed distributions could continue learning across the full set of boosting iterations albeit marginally and early stopping never occurs. Infrequently occurring tail classes — even those with fewer than 100 samples — could still find useful splits hundreds if not thousands of rounds in. The head classes converged long ago — micro recall (dominated by head classes) stabilizes early, but macro recall (averaged equally across all classes) continues to improve as tail classes learn. Training cannot stop without sacrificing the accuracy of tail classes and training runs to the fully configured number of boosting iterations.
With 2000 bins, 38 features, and 400 classes, the histogram construction step — the same inner loop from binary classification — now evaluates 30+ million independent accumulation targets per node split. Each operation is still just an addition, but there are millions of them with zero dependencies between them, repeated every iteration for 2000 iterations.
The workload has shifted from narrow serial tree logic to massively parallel streaming reduction operations.
Workload Fits GPU Acceleration
A GPU executes thousands of things simultaneously rather than executing one thing fast. Each CUDA core is individually slower than a CPU core. The advantage only shows up when there’s enough independent work to fill all cores every cycle.
For the binary classification profile described earlier — 12,800 accumulation targets per split — a GPU with thousands of cores has no advantage. Most cores sit idle with nothing to process. But the multi-class workload with 30 million independent accumulations per split saturates every available core. Combined with GPU memory bandwidth of 600–900 GB/s (the A10G/V100 tier where XGBoost lives economically) versus CPU’s ~50 GB/s, the expanded gradient array streams through 12–18× faster. Higher-bandwidth GPUs exist (A100 at ~2 TB/s) but XGBoost can’t utilize the additional bandwidth — histogram accumulation is bottlenecked by atomic additions into bin accumulators, not by data arrival rate. Faster memory buses don’t translate to faster boosting rounds. This coincides with findings by Mitchell and Frank (2017), who demonstrated significant speedups from GPU memory bandwidth on histogram construction, the same bottleneck operation described here.
The high class count isn’t just a memory problem — it’s what creates the parallelism that justifies GPU use and saturates the many small GPU cores.
Wider parallelism maps to GPU core utilization, but the same properties that make the workload wide — high class count, aggressive binning, large data — also expand its memory footprint. We now need to ensure we have the VRAM to fit the needs of this new problem. Unlike CPU nodes where RAM can be provisioned into the terabytes cheaply and on demand, VRAM is a fixed physical constraint per card — 24 GB, 40 GB, 48 GB, up to 192 GB on the latest hardware. More importantly, GPU instances with sufficient VRAM are harder to provision on demand and more expensive per GB than general-purpose RAM. Overshooting on CPU means a slightly larger node at marginal cost. Overshooting on GPU means reserving scarce, premium hardware. Right-sizing VRAM is not optional — it requires upfront calculation. Understanding exactly what must reside in VRAM — and sizing the cluster so that it fits — is the remaining engineering problem.
[Image 1: Architecture comparison — Left: CPU cluster with large RAM nodes, small CPUs mostly idle, gradient array in shared host memory. Right: GPU cluster with gradient array in dedicated VRAM, CUDA cores fully saturated.]
VRAM Anatomy
GPU memory during XGBoost training holds three components. Understanding them is the entire sizing exercise.
1. The ELLPACK DMatrix
The training data lives in VRAM for the entire duration of training — every boosting round reads from it. On GPU, XGBoost stores data in ELLPACK format: a compressed structure where raw feature values have been quantised into integer bin indices. Each cell is a bin index rather than a float.
What affects DMatrix size:
- Rows per GPU — scales linearly. More data per shard, larger DMatrix.
- Number of features — scales linearly.
- max_bin — determines the bit-width of bin indices. Higher max_bin means more bits per cell.
- Density — sparse data compresses; dense data (common in tabular ML) stores the full grid.
In our testing, dense data with max_bin=2000 and 38 features at 5.5M rows per GPU shard occupied roughly 13 GB in ELLPACK form. This never leaves VRAM. It's read thousands of times across boosting rounds but never modified.
2. NCCL Communication Buffers
Distributed training requires AllReduce to synchronise histogram statistics across GPUs after each round. NCCL allocates buffers for this coordination.
What affects NCCL overhead:
- Number of workers — more participants means larger collective buffers.
- Histogram payload size —
max_bin × features × classes × 2 × 4 bytesis the data being synced each round.
For typical configurations, this is modest — hundreds of MB, not GB. But with many small GPUs (e.g., 30+ single-GPU nodes) the coordination overhead grows relative to compute. Fewer, larger-VRAM GPUs reduce this cost per round.
3. Gradient and Hessian Arrays
This is the cost of admission for boosting. Before the first iteration can begin, XGBoost allocates: rows_per_gpu × num_classes × 2 × 4 bytes. The factor of 2 is the first derivative (gradient) and second derivative (hessian) — both needed to compute optimal leaf weights via Newton's method.
Get Garrett Frere’s stories in your inbox
Join Medium for free to get updates from this writer.
What affects gradient/hessian size:
- Rows per GPU — scales linearly.
- Number of target classes — the multiplier that distinguishes trivial workloads from VRAM-intensive ones. Binary classification: 2 values per sample. 420-class softprob: 840 values per sample. Same rows, same features, 420× the gradient memory.
For 420 classes and 5.5M rows per GPU: 5.5M × 420 × 2 × 4 bytes ≈ 18-20 GB.
What gets overwritten vs. what grows
The gradient/hessian arrays are overwritten in place every round. After each tree is added, predictions update, and new gradients are calculated against the current ensemble’s residuals. Round 1’s allocation is the same allocation at round 2000. The VRAM footprint of boosting is essentially static — no accumulation, no growing buffers, no surprise allocations mid-training.
The only growth comes from the tree structures themselves (leaf weights, split metadata) and histogram working memory during split evaluation. In our testing, this additional overhead was minimal — approximately 1–2 GB beyond the baseline of DMatrix + gradients across a full 2000-round run.
Sizing from first principles
Total VRAM per GPU = DMatrix + NCCL buffers + Gradients/Hessians + ~2 GB overhead
The calculation is deterministic. Before running anything, you know:
- Your data volume and how you’ll shard it (rows per GPU)
- Your feature count and max_bin (DMatrix size)
- Your class count (gradient array size)
If the sum exceeds the GPU’s VRAM, it will fail on the first iteration — not gradually, not probabilistically, but immediately. Size your cluster so that each GPU has headroom above this sum, and the job will run to completion regardless of how many boosting rounds you configure.
QuantileDMatrix: The Loading Optimization
With VRAM constraints understood, the question becomes: how does data get into VRAM efficiently?
QuantileDMatrix in XGBoost 2.0+ provides streaming sketch-based data construction. Rather than loading all raw data into GPU memory then building the ELLPACK structure, it streams batches, sketches quantile boundaries, maps values to integer bin indices, and discards the raw floats. Peak memory during loading never holds full raw data and full internal structure simultaneously.
No precision trade-off. The hist method already quantises internally — it only evaluates splits at bin boundaries regardless. QuantileDMatrix performs the identical quantisation step earlier and more memory-efficiently. The resulting model is the same.
The distinction matters: QuantileDMatrix reduces peak memory during data loading, not during training. Once the ELLPACK page is built, it occupies the same space regardless of how it was constructed. The benefit is that datasets which would OOM during DMatrix construction now load successfully — widening the range of problems that fit on a given GPU.
The XGBEstimator distributor uses QuantileDMatrix automatically on each worker — no code change needed. This is available in Snowflake Container Runtime v2.5 and later (which ships XGBoost 3.2.0).
Diagnosing VRAM OOM
When sizing is wrong, the failure is immediate and unambiguous. It always occurs on the first boosting iteration — not mid-training, not after hours of progress. The sequence is predictable:
First, data loads into VRAM as a compressed ELLPACK DMatrix. The worker logs show Generating new Ellpack page followed by Ellpack is dense. This succeeds — the DMatrix fits.
Then iteration zero fires. XGBoosterUpdateOneIter allocates the gradient/hessian array on top of the DMatrix already in memory. The CUDA allocator fails and XGBoost reports the exact state:
Memory allocation error on worker 0: std::bad_alloc: cudaErrorMemoryAllocation: out of memory
- Free memory: 10.95GB
- Requested memory: 20.89GBFree memory is what remained after the DMatrix loaded. Requested is the gradient array for the first iteration. The gap is the deficit — in this case, roughly 10 GB short. Not a mystery, not a race condition, not a driver bug. Arithmetic.
In a Ray distributed setup, every rank fails simultaneously with the same error at the same call site. This distinguishes VRAM sizing from network failures or data corruption — those produce asymmetric errors across ranks.
The fix follows directly from the numbers. Either add GPUs to reduce rows per shard (reducing both the DMatrix footprint and the gradient array proportionally), or move to GPUs with more VRAM. You know the deficit before you re-run.
Practicalities
For teams on Snowflake Container Runtime, the transition is a configuration change:
from snowflake.ml.modeling.distributors.xgboost import XGBEstimator, XGBScalingConfig
estimator = XGBEstimator(
params=params,
scaling_config=XGBScalingConfig(use_gpu=True)
)
booster = estimator.fit(
dataset=train_connector,
input_cols=input_cols,
label_col=label_col,
eval_set=eval_connector,
verbose_eval=100
)Same API, same data source, same hyperparameters, same model output. Snowflake automatically discovers and uses all available GPU resources on the cluster — no manual worker count or device assignment needed.
XGBoost’s GPU requirements are modest by modern standards. It doesn’t need high-speed interconnects (HSI) between nodes — standard PCIe is sufficient since coordination between workers is limited to small AllReduce payloads, not full model synchronization. It doesn’t use tensor cores. Older architectures — Ampere and earlier — are perfectly adequate. The GPU capabilities XGBoost actually exercises have been widely available for years, and that hardware keeps getting cheaper.
When to Stay on CPU
XGBoost training that remains narrow can stay on CPU. The workload is serial and iterative — CPU-appropriate when target class count is low, default max_bin is sufficient, early stopping triggers within a few hundred rounds, and data fits in worker memory without over-provisioning. This describes many XGBoost workloads, but as data becomes more granular and classification targets more complex, multi-class problems naturally emerge and deviate from the traditional binary profile.
The crossover occurs when high cardinality, aggressive binning, and sustained boosting compound — transforming the inner loop from a CPU-tractable operation into wide parallel accumulation. When that happens, GPU is the architecturally correct choice.
With consumption-based billing models that many cloud computing platforms employ, architecturally correct becomes economically correct — the faster you complete training, the sooner you release resources and stop consuming compute. A job that finishes in hours on GPU rather than days on CPU costs less in total even at a higher per-hour rate. At this crossover point, the GPU option is the economical one.
References
- Mitchell R, Frank E. (2017) “Accelerating the XGBoost algorithm using GPU computing.” PeerJ Computer Science 3:e127
- XGBoost GPU Support: https://xgboost.readthedocs.io/en/latest/gpu/index.html
- XGBoost QuantileDMatrix: https://xgboost.readthedocs.io/en/stable/python/python_api.html#xgboost.QuantileDMatrix
- Snowflake Container Runtime GPU 2.6: https://docs.snowflake.com/en/developer-guide/snowflake-ml/container-runtime/releases/gpu/2_6
