[Python] VLLM Objects Cause Memory Errors When Created in a Function even when explicitly...

I'm encountering an issue when using the VLLM library in Python. Specifically, when I create a VLLM model object inside a function, I run into memory problems and cannot clear the GPU memory effectively, even after deleting objects and using torch.cuda.empty_cache().

The problem occurs when I try to instantiate a LLM object inside a function, but it does not happen if I instantiate the object in the parent process or global scope. This suggests that VLLM has issues with creating and managing objects in functions, which leads to memory retention and GPU exhaustion.

Here’s a simplified version of the code:

import torch
import gc
from vllm import LLM

def run_vllm_eval(model_name, sampling_params, path_2_eval_dataset):
# Instantiate LLM in a function
llm = LLM(model=model_name, dtype=torch.float16, trust_remote_code=True)

# Run some VLLM inference or evaluation here (simplified)
result = llm.generate([path_2_eval_dataset], sampling_params)

# Clean up after inference
del llm
gc.collect()
torch.cuda.empty_cache()

# After this, GPU memory is not cleared properly and causes OOM errors
run_vllm_eval()
run_vllm_eval()
run_vllm_eval()


but

llm = run_vllm_eval2()
llm = run_vllm_eval2(llm)
llm = run_vllm_eval2(llm)


Works.

Even after explicitly deleting the LLM object and clearing the cache, the GPU memory is not properly freed, leading to out-of-memory (OOM) errors when trying to load or run another model in the same script.

Things I've Tried:

• Deleting the LLM object with del.
• Running gc.collect() to trigger Python's garbage collection.
• Using torch.cuda.empty_cache() to clear CUDA memory.
• Ensuring no VLLM objects are instantiated in the parent process.

None of these seem to fix the issue when the LLM object is created within a function.

Questions:

• Has anyone encountered similar memory issues when creating VLLM objects inside functions?
• Is there a recommended way to manage or clear VLLM objects in a function to prevent GPU memory retention?
• Are there specific VLLM handling techniques that differ from standard Hugging Face or PyTorch models in this context?

Continue reading...