Possible memory leak when using XGBoost in conjunction with Torch
- 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