XGBoost `Booster` Ignores `device=cpu` When Loading Model from GPU Training
- Dominant language
- C++
- Stars
- 28.8k
- Forks
- 8.9k
- Avg merge
- 1d 12h
- Merged PRs (30d)
- 54
Description
**Description:**
When training a model on GPU and then loading it on CPU using `xgb.Booster`, the `device` parameter appears to be set correctly in `save_config()`, but inference still utilizes the GPU unexpectedly. This results in high GPU memory usage and slow predictions.
I am using a Train Loop Fn for Out of Sample predictions, but the problem allready occurs in the first itteration of the loop.
---
### **Environment:**
- **XGBoost Version:** 2.1.3
- **Python Version:** 12.8
- **OS:** Windows 11
---
### **Reproduction Code:**
```python
import xgboost as xgb
import numpy as np
import polars as pl
import gc
from typing import Optional
# Placeholder values
FEATURES = ["feature1", "feature2", "feature3"] # Example feature columns
TARGETS = "target" # Example target column
WEIGHTS = "weight" # Example weight column
train_df = pl.DataFrame()
valid_df = pl.DataFrame()
eval_df = pl.DataFrame()
def create_xgb_matricies(train_df: pl.DataFrame, valid_df: pl.DataFrame, eval_df: pl.DataFrame):
X_train = train_df[FEATURES].to_numpy()
y_train = train_df[TARGETS].to_numpy()
w_train = train_df[WEIGHTS].to_numpy()
dtrain = xgb.DMatrix(X_train, label=y_train, weight=w_train)
del train_df, X_train, y_train, w_train
X_valid = valid_df[FEATURES].to_numpy()
y_valid = valid_df[TARGETS].to_numpy()
w_valid = valid_df[WEIGHTS].to_numpy()
dvalid = xgb.DMatrix(X_valid, label=y_valid, weight=w_valid)
del valid_df, X_valid, y_valid, w_valid
X_eval = eval_df[FEATURES].to_numpy()
y_eval = eval_df[TARGETS].to_numpy()
w_eval = eval_df[WEIGHTS].to_numpy()
deval = xgb.DMatrix(X_eval, label=y_eval, weight=w_eval)
del eval_df, X_eval, y_eval, w_eval
gc.collect()
return dtrain, dvalid, deval
def get_pred(train: xgb.DMatrix, valid: xgb.DMatrix, eval: xgb.DMatrix, params):
gpu_model = xgb.train(params, train, num_boost_round=1000,
evals=[(valid, 'val')], early_stopping_rounds=10, verbose_eval=False)
print("training, done!")
# 🔥 Save the model
gpu_model.save_model("gpu_model.json")
print(gpu_model.save_config()) # Should show 'device'='cpu'
del gpu_model
cpu_model = xgb.Booster(params={"device": "cpu"})
cpu_model.load_model('gpu_model.json')
print(cpu_model.save_config()) # Should show 'device'='cpu'
print("loading done!")
pred = cpu_model.predict(eval)
del cpu_model
return pred
def get_oos_probas(mats: list[xgb.DMatrix], params=None):
params = params or {'objective': 'binary:logistic', 'eval_metric': 'logloss', 'seed': 42, 'device': 'cuda', 'learning_rate': 0.05}
oos_preds: list[Optional[np.ndarray]] = [None] * len(mats)
for i in range(len(mats)):
train_idx, val_idx, test_idx = i, (i+1)%3, (i+2)%3
oos_preds[test_idx] = get_pred(mats[train_idx], mats[val_idx], mats[test_idx], params)
return oos_preds
dtrain, dvalid, deval = create_xgb_matricies(train_df, valid_df, eval_df)
oos_preds = get_oos_probas([dtrain, dvalid, deval]) # , PARAMS
```
---
### **Expected Behavior:**
- After loading the model with `device=cpu`, inference should be performed entirely on the CPU.
- Task Manager should show CPU utilization during inference.
### **Observed Behavior:**
- Despite `save_config()` showing `device=cpu`, inference still uses GPU resources.
- GPU utilization is **0% compute** but **maxed-out GPU memory** in Task Manager.
- This results in **slow predictions**.
---
### **Additional Notes:**
- I was not able to do predictions on CPU unless training on the CPU.
- All workarounds tried where not effective.
**Would appreciate any guidance or fixes for this issue!**
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.