概览
What is WarpGBM?
WarpGBM is a high-performance, GPU-accelerated Gradient Boosted Decision Tree (GBDT) library built with PyTorch and custom CUDA kernels. It supports regression, binary, and multiclass classification and includes an invariant learning algorithm (DES) that learns signals stable across shifting distributions. It is for data scientists and machine learning engineers who need speed and distribution-robust models.
How to use WarpGBM?
Install via pip install warpgbm or pip install git+https://github.com/jefferythewind/warpgbm.git. Requires PyTorch with CUDA support. Use the Scikit-Learn-compatible API: create a WarpGBM object with parameters like objective, max_depth, n_estimators, call fit(X, y) and predict(X). For invariant learning, pass an era_id array to fit.
Key features of WarpGBM
- GPU-native CUDA kernels for binning, histograms, splits, and inference
- Unified infrastructure for regression, binary, and multiclass classification
- Invariant learning via Directional Era-Splitting (DES) algorithm
- Scikit-Learn compatible – drop‑in replacement
- Warm start and checkpointing for incremental training
- Pre-binned data optimization and mixed precision (float32/int8) support
Use cases of WarpGBM
- Training high-speed gradient boosting models on large datasets (1M+ rows)
- Building models that generalize across shifting market regimes or time periods
- Multiclass classification with early stopping and rich evaluation metrics
- Incremental model development and production updates without retraining from scratch
- Hyperparameter tuning with checkpointing to decide on the number of trees
FAQ from WarpGBM
How does WarpGBM compare to XGBoost, LightGBM, and CatBoost?
On synthetic data (1M rows × 1K features, Google Colab L4 GPU), WarpGBM trains 2× faster than XGBoost and 23× faster than CatBoost while achieving comparable correlation. LightGBM inference is faster, but WarpGBM trains similarly to LightGBM and beats it on pre-binned Numerai data (13× faster).
What are the prerequisites for installing WarpGBM?
PyTorch with CUDA support is required. The library uses custom CUDA kernels; a GPU is strongly recommended, though CPU mode (device='cpu') is available.
Can I train a model without a GPU?
Yes, set device='cpu' in the constructor. However, performance will be significantly slower – GPU acceleration is a core feature.
What is the invariant learning feature (DES)?
When you pass an era_id array (marking which time period/regime
