XGBoost Distributed Training · Multi-Node and Dask Integration

XGBoost supports distributed training across multiple machines using its native distributed framework or through cluster computing systems (Dask, Spark, Ray). This article covers the key approaches and their trade-offs.

Native Distributed Framework

XGBoost uses the Rabit communication library for all-reduce operations (gradient aggregation across workers). Architecture:

  1. One master node coordinates the training
  2. Multiple worker nodes hold data partitions
  3. Workers compute local gradient histograms, Rabit performs all-reduce to aggregate globally
# On master
xgboost params.conf model_out=model.json num_workers=4

# On each worker (worker_index 0,1,2,3)
xgboost params.conf model_in=model.json

Rabit handles fault tolerance: if a worker fails, recovery from checkpoint is supported.

Dask Integration

Dask-XGBoost is the most practical distributed training API:

from dask.distributed import Client
import xgboost as xgb
import dask.dataframe as dd

client = Client('scheduler:8786')
df = dd.read_csv('s3://bucket/data-*.csv')

# Dask automatically partitions data across workers
model = xgb.dask.DaskXGBRegressor(
    n_estimators=100,
    max_depth=6,
    tree_method='hist'  # hist required for distributed
)
model.fit(df[features], df[target])

Key advantage: Dask handles data partitioning, worker management, and fault tolerance. The XGBoost Dask API mirrors scikit-learn.

When to Go Distributed

xgboost-org 配图

Single node sufficient when: data < 100M rows, training time < 1 hour on a single machine.

Distributed needed when: data > 100M rows, or training time becomes a bottleneck (want sub-10-minute training on 500M rows).

Memory constraint: 100M rows × 50 features × 8 bytes = 40 GB. A single node with 128 GB RAM handles this. Distribution becomes necessary around 500M+ rows.

Practical Tips

xgboost-org 配图

  • Use tree_method='hist' (required for distributed)
  • Set max_bin=256 (default, balances memory and accuracy)
  • Partition data evenly across workers (10–50 partitions per worker)
  • For cloud: 8–16 vCPUs per worker, enough RAM to hold 3× the partition size (for intermediate buffers)