dmlc / dmlc/xgboost

XGBoost `Booster` Ignores `device=cpu` When Loading Model from GPU Training

Open
#11,199 12 comments 0 reactions 0 assignees View on GitHub
status: need update
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.