Check if XGBoost Supports GPU

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 to True. If present, it indicates that XGBoost was built with CUDA support, enabling GPU acceleration.
If the USE_CUDA field is missing or shows False, 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.

See Also