WarpGBM is a high-performance, GPU-accelerated Gradient Boosted Decision Tree (GBDT) library built with PyTorch and CUDA. It offers blazing-fast histogram-based training and efficient prediction, with compatibility for research and production workflows.
- GPU-accelerated training and histogram construction using custom CUDA kernels
- Drop-in scikit-learn style interface
- Supports pre-binned data or automatic quantile binning
- Simple install with
pip
In this benchmark we compare the speed and in-sample correlation of WarpGBM v0.1.21 against LightGBM, XGBoost and CatBoost, all with their GPU-enabled versions. This benchmark runs on Google Colab with the L4 GPU environment.
WarpGBM: corr = 0.8882, train = 18.7s, infer = 4.9s
XGBoost: corr = 0.8877, train = 33.1s, infer = 8.1s
LightGBM: corr = 0.8604, train = 30.3s, infer = 1.4s
CatBoost: corr = 0.8935, train = 400.0s, infer = 382.6s
Colab Notebook: https://colab.research.google.com/drive/16U1kbYlD5HibGbnF5NGsjChZ1p1IA2pK?usp=sharing
pip install git+https://github.com/jefferythewind/warpgbm.gitThis installs the latest version directly from GitHub and compiles CUDA extensions on your machine using your local PyTorch and CUDA setup. It's the most reliable method for ensuring compatibility and staying up to date with the latest features.
pip install warpgbmThis installs from PyPI and also compiles CUDA code locally during installation. This method works well if your environment already has PyTorch with GPU support installed and configured.
Tip:
If you encounter an error related to mismatched or missing CUDA versions, try installing with the following flag. This is currently required in the Colab environments.pip install warpgbm --no-build-isolation
Thank you, ShatteredX, for providing working instructions for a Windows installation.
git clone https://github.com/jefferythewind/warpgbm.git
cd warpgbm
python setup.py bdist_wheel
pip install .\dist\warpgbm-0.1.15-cp310-cp310-win_amd64.whl
Before either method, make sure you’ve installed PyTorch with GPU support:
https://pytorch.org/get-started/locally/
import numpy as np
from sklearn.datasets import make_regression
from time import time
import lightgbm as lgb
from warpgbm import WarpGBM
# Create synthetic regression dataset
X, y = make_regression(n_samples=100_000, n_features=500, noise=0.1, random_state=42)
X = X.astype(np.float32)
y = y.astype(np.float32)
# Train LightGBM
start = time()
lgb_model = lgb.LGBMRegressor(max_depth=5, n_estimators=100, learning_rate=0.01, max_bin=7)
lgb_model.fit(X, y)
lgb_time = time() - start
lgb_preds = lgb_model.predict(X)
# Train WarpGBM
start = time()
wgbm_model = WarpGBM(max_depth=5, n_estimators=100, learning_rate=0.01, num_bins=7)
wgbm_model.fit(X, y)
wgbm_time = time() - start
wgbm_preds = wgbm_model.predict(X)
# Results
print(f"LightGBM: corr = {np.corrcoef(lgb_preds, y)[0,1]:.4f}, time = {lgb_time:.2f}s")
print(f"WarpGBM: corr = {np.corrcoef(wgbm_preds, y)[0,1]:.4f}, time = {wgbm_time:.2f}s")LightGBM: corr = 0.8742, time = 37.33s
WarpGBM: corr = 0.8621, time = 5.40s
WarpGBM can save additional training time if your dataset is already pre-binned. The Numerai tournament data is a great example:
import pandas as pd
from numerapi import NumerAPI
from time import time
import lightgbm as lgb
from warpgbm import WarpGBM
import numpy as np
napi = NumerAPI()
napi.download_dataset('v5.0/train.parquet', 'train.parquet')
train = pd.read_parquet('train.parquet')
feature_set = [f for f in train.columns if 'feature' in f]
target = 'target_cyrus'
X_np = train[feature_set].astype('int8').values
Y_np = train[target].values
# LightGBM
start = time()
lgb_model = lgb.LGBMRegressor(max_depth=5, n_estimators=100, learning_rate=0.01, max_bin=7)
lgb_model.fit(X_np, Y_np)
lgb_time = time() - start
lgb_preds = lgb_model.predict(X_np)
# WarpGBM
start = time()
wgbm_model = WarpGBM(max_depth=5, n_estimators=100, learning_rate=0.01, num_bins=7)
wgbm_model.fit(X_np, Y_np)
wgbm_time = time() - start
wgbm_preds = wgbm_model.predict(X_np)
# Results
print(f"LightGBM: corr = {np.corrcoef(lgb_preds, Y_np)[0,1]:.4f}, time = {lgb_time:.2f}s")
print(f"WarpGBM: corr = {np.corrcoef(wgbm_preds, Y_np)[0,1]:.4f}, time = {wgbm_time:.2f}s")Results (Google Colab Pro, A100 GPU):
LightGBM: corr = 0.0703, time = 643.88s
WarpGBM: corr = 0.0660, time = 49.16s
You can try WarpGBM in a live Colab notebook using real pre-binned Numerai tournament data:
No installation required — just press "Open in Playground", then Run All!
num_bins: Number of histogram bins to use (default: 10)max_depth: Maximum depth of trees (default: 3)learning_rate: Shrinkage rate applied to leaf outputs (default: 0.1)n_estimators: Number of boosting iterations (default: 100)min_child_weight: Minimum sum of instance weight needed in a child (default: 20)min_split_gain: Minimum loss reduction required to make a further partition (default: 0.0)histogram_computer: Choice of histogram kernel ('hist1','hist2','hist3') (default:'hist3')threads_per_block: CUDA threads per block (default: 32)rows_per_thread: Number of training rows processed per thread (default: 4)L2_reg: L2 regularizer (default: 1e-6)colsample_bytree: Proportion of features to subsample to grow each tree (default: 1)
.fit(
X, # numpy array (float or int) 2 dimensions (num_samples, num_features)
y, # numpy array (float or int) 1 dimension (num_samples)
era_id=None, # numpy array (int) 1 dimension (num_samples)
X_eval=None, # numpy array (float or int) 2 dimensions (eval_num_samples, num_features)
y_eval=None, # numpy array (float or int) 1 dimension (eval_num_samples)
eval_every_n_trees=None, # const (int) >= 1
early_stopping_rounds=None, # const (int) >= 1
eval_metric='mse' # string, one of 'mse' or 'corr'. For corr, loss is 1 - correlation(y_true, preds)
)
Train with optional validation set and early stopping.
.predict(
X # numpy array (float or int) 2 dimensions (predict_num_samples, num_features)
)
Predict on new data, using parallelized CUDA kernel.
WarpGBM builds on the shoulders of PyTorch, scikit-learn, LightGBM, and the CUDA ecosystem. Thanks to all contributors in the GBDT research and engineering space.
- Vectorized predict function replaced with CUDA kernel (
warpgbm/cuda/predict.cu), parallelizing per sample, per tree.
- Adjust gain in split kernel and added support for an eval set with early stopping based on MSE.
- Added
colsample_bytreeparameter and new test using Numerai data.
- Fix Memory bugs in prediction and colsample bytree logic. Added "corr" eval metric.