XGBoost is known for its excellent performance, and using GPU acceleration can further speed up the training process, especially for large datasets.
XGBoost provides GPU support, but it depends on the installation and system configuration. Here’s how you can check if your XGBoost installation supports GPU.
import xgboost
print(xgboost.build_info())
Output (your output may differ):
{'BUILTIN_PREFETCH_PRESENT': True, 'CLANG_VERSION': [14, 0, 0], 'DEBUG': False, 'MM_PREFETCH_PRESENT': True, 'USE_CUDA': False, 'USE_FEDERATED': False, 'USE_NCCL': False, 'USE_OPENMP': True, 'USE_RMM': False, 'libxgboost': '/usr/local/lib/python3.12/site-packages/xgboost/lib/libxgboost.dylib'}
We can make the output easier to read with a little formatting, for example:
import xgboost
build_info = xgboost.build_info()
for name in sorted(build_info.keys()):
print(f'{name}: {build_info[name]}')
Output (your output may differ):
BUILTIN_PREFETCH_PRESENT: True
CLANG_VERSION: [14, 0, 0]
DEBUG: False
MM_PREFETCH_PRESENT: True
USE_CUDA: False
USE_FEDERATED: False
USE_NCCL: False
USE_OPENMP: True
USE_RMM: False
libxgboost: /usr/local/lib/python3.12/site-packages/xgboost/lib/libxgboost.dylib
Here’s what to look for:
- Check if the output includes a
USE_CUDA
field set toTrue
. If present, it indicates that XGBoost was built with CUDA support, enabling GPU acceleration. - If the
USE_CUDA
field is missing or showsFalse
, it means XGBoost was built without GPU support.
To use XGBoost with GPU acceleration, you need:
- An NVIDIA GPU with CUDA compute capability.
- CUDA toolkit installed on your system.
- XGBoost installed with the appropriate GPU libraries.
Installing XGBoost with GPU support depends on your system and package manager.
Here are a few common methods:
- Using pip:
pip install xgboost-gpu
- Using conda:
conda install -c conda-forge xgboost-gpu
- Building from source with GPU support enabled
If you don’t see GPU support in your current XGBoost installation, you may need to reinstall it with the GPU-enabled version.
Note that even with GPU support, XGBoost will only use the GPU if the dataset and model are large enough to benefit from GPU acceleration. For smaller datasets, the overhead of copying data to and from the GPU may outweigh the performance gains.
To explicitly enable or disable GPU usage, you can set the device
parameter to cuda
or gpu
when training your model:
# Enable GPU acceleration
model = xgb.XGBClassifier(device='cuda')
# Disable GPU acceleration (use CPU)
model = xgb.XGBClassifier(device='cpu')
By checking for GPU support and configuring XGBoost accordingly, you can leverage the power of your GPU to accelerate model training and improve overall performance.