dmlc / dmlc/xgboost

Pred_contribs in 2.1.1 takes significantly more gpu memory than in 1.4.2

Open
#10,936 6 comments 2 reactions 0 assignees View on GitHub
Dominant language
C++
Stars
28.8k
Forks
8.9k
Avg merge
1d 12h
Merged PRs (30d)
54

Description

I've noticed that using pred_contribs to generate shap values takes significantly more gpu memory in XGBoost 2.1.1 vs 1.4.2.
This can lead to having issues with generating shap values, where no issue was previously present.

GPU memory comparison:
1.4.2 - 3090
1.7.6 - 4214
2.1.1 - 5366

Short example used to demonstrate:
```
from typing import Tuple
import subprocess
import pandas as pd
from ucimlrepo import fetch_ucirepo
from sklearn.model_selection import train_test_split
from xgboost import XGBClassifier, DMatrix

def download_data() -> Tuple[pd.DataFrame, pd.DataFrame]:
# fetch dataset
diabetes_binary = fetch_ucirepo(id=891)
# data (as pandas dataframes)
X = diabetes_binary.data.features
y = diabetes_binary.data.targets
return X, y


def prep_dataset(X: pd.DataFrame, y: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]:
# split dataset into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=0)

return X_train, X_test, y_train, y_test

def train_model(X_train: pd.DataFrame, y_train: pd.DataFrame) -> XGBClassifier:
# train a model
xgb_params = {
"objective": "binary:logistic",
"n_estimators": 2000,
"max_depth": 13,
"learning_rate": 0.1,
"tree_method": "gpu_hist",
}
model = XGBClassifier(**xgb_params)
model.fit(X_train, y_train["Diabetes_binary"])
return model

def call_shap_log_usage(model: XGBClassifier, test_data: pd.DataFrame) -> None:
output_file = open("output.txt", mode='w')
proc = subprocess.Popen('./log_usage.sh', stdout=output_file, stderr=subprocess.STDOUT)

booster = model.get_booster()
booster.set_param({"predictor": "gpu_predictor"})
dmatrix = DMatrix(test_data)
shap_values = booster.predict(dmatrix, pred_contribs=True)

proc.terminate()

return shap_values

if __name__ == '__main__':
X, y = download_data()
X_train, X_test, y_train, y_test = prep_dataset(X, y)
model = train_model(X_train, y_train)
call_shap_log_usage(model, X_test)
```

with the following bash script used for generating memory usage:
```
#!/bin/bash
a=0
while true; do
b=$(nvidia-smi --query-gpu=memory.used --format=csv|grep -v memory|awk '{print $1}')
[ $b -gt $a ] && a=$b && echo $a
sleep .5
done
```

All tests run on Ubuntu 20.04.6 LTS.
Requirements with only the xgb version (and the device/tree method parameters) being changed between tests:
```
pandas==1.3.5
numpy==1.22.4
ipykernel==5.5.6
xgboost==1.4.2
scikit-learn==1.3.2
ucimlrepo==0.0.7
```

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.