dmlc / dmlc/xgboost

Possible memory leak when using XGBoost in conjunction with Torch

Open
#9,293 1 comment 0 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 had an issue in one of the services I work on, where it would use more and more memory until crashing. After some digging around I was able to reduce it to the following script:
```python
import argparse
import logging
import math
import os

import psutil
import torch
import torchvision
import xgboost

import numpy as np

def main() -> None:
process = psutil.Process(os.getpid())

parser = argparse.ArgumentParser()
parser.add_argument("xgboost_model_path")
args = parser.parse_args()

feature_extractor = torchvision.models.vit_b_16(num_classes=512)

predictor = xgboost.XGBClassifier(
base_score=0.5,
booster=None,
colsample_bylevel=1,
colsample_bynode=1,
colsample_bytree=1,
gamma=0,
gpu_id=-1,
importance_type="gain",
interaction_constraints=None,
learning_rate=0.3,
max_delta_step=0,
max_depth=10,
min_child_weight=1,
missing=math.nan,
monotone_constraints=None,
n_estimators=300,
n_jobs=32,
num_parallel_tree=1,
objective="multi:softprob",
random_state=0,
reg_alpha=0,
reg_lambda=1,
scale_pos_weight=None,
subsample=1,
tree_method=None,
validate_parameters=False,
verbosity=0,
)
predictor.load_model(args.xgboost_model_path)

frames_torch = torch.rand((1, 3, 224, 224), device="cpu")
i = 0
while True:
with torch.no_grad():
embedding = feature_extractor(frames_torch).numpy().mean(axis=0)
if i == 0:
logging.warning(f"Mem usage (embedding) {process.memory_percent()}")
features = np.expand_dims(embedding, 0)
predictor.predict_proba(features)
i = (i + 1) % 10

if __name__ == "__main__":
main()
```
which uses the following xgboost model: [xgboost_classifier.txt](https://github.com/dmlc/xgboost/files/11720066/xgboost_classifier.txt) (txt format becuase github doesn't allow JSON, apparently).

I don't know if the memory really is leaked, but at the very least its usage grows without bounds. I left this script running for a day and memory usage grew from 640MB to about 9GB. What's interesting is that this seems to depend on the import order, if XGBoost is imported before torch and torchvision the issue doesn't reproduce (I didn't leave it to run for the same amount of time but I didn't see an upwards trend that's clearly visible otherwise).

I use python 3.9.15. Here are the installed packages:
```
certifi==2023.5.7
charset-normalizer==3.1.0
cmake==3.26.4
filelock==3.12.1
idna==3.4
Jinja2==3.1.2
joblib==1.2.0
lit==16.0.5.post0
MarkupSafe==2.1.3
mpmath==1.3.0
networkx==3.1
numpy==1.24.3
nvidia-cublas-cu11==11.10.3.66
nvidia-cuda-cupti-cu11==11.7.101
nvidia-cuda-nvrtc-cu11==11.7.99
nvidia-cuda-runtime-cu11==11.7.99
nvidia-cudnn-cu11==8.5.0.96
nvidia-cufft-cu11==10.9.0.58
nvidia-curand-cu11==10.2.10.91
nvidia-cusolver-cu11==11.4.0.1
nvidia-cusparse-cu11==11.7.4.91
nvidia-nccl-cu11==2.14.3
nvidia-nvtx-cu11==11.7.91
Pillow==9.5.0
psutil==5.9.5
requests==2.31.0
scikit-learn==1.2.2
scipy==1.10.1
sympy==1.12
threadpoolctl==3.1.0
torch==2.0.1
torchvision==0.15.2
triton==2.0.0
typing_extensions==4.6.3
urllib3==2.0.3
xgboost==1.7.5
```
Here's the system I'm using:
```
$ uname -rv
5.15.0-1036-aws #40~20.04.1-Ubuntu SMP Mon Apr 24 00:21:13 UTC 2023
```
I don't know if this is an issue with XGBoost specifically, but hopefully, someone here has ideas on what could be the cause. I suspect that this might be related to some library that both Torch and XGBoost use, like OpenMP.

Contributor guide

No contributing guide indexed for this repository

Research direction

The only reproducer is the inline Python script, using torch/torchvision inference followed by XGBClassifier.predict_proba and the linked xgboost_classifier.txt model. Run it with the listed versions, compare import orders and memory over time, then determine whether growth is in XGBoost, Torch, or a shared dependency; done means a confirmed cause and verified bounded memory use.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python, pytorch
Domain
machine-learning, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.