[Usage]: How to release GPU of vLLM model in python code

[quanshr]

Your current environment

llm = LLM(model=model1_path, tensor_parallel_size=torch.cuda.device_count())
llm = LLM(model=model2_path, tensor_parallel_size=torch.cuda.device_count())

It will cause CUDA out of memory when execute the second line.

How would you like to use vllm

I want to use two model in pipeline in one python code to infer. When finish inference on the first model, how to release this model and release GPU memory to load another one, since directly reloading may cause CUDA OUT OF MEMORY for it doesn't release the first one.

[DarkLight1337]

Please see #5716. For now, you have to do it manually as follows (based on the code in tests/conftest.py:

import contextlib
import gc

import torch
from vllm.distributed import (destroy_distributed_environment,
                              destroy_model_parallel)
from vllm.utils import is_cpu

def cleanup():
    destroy_model_parallel()
    destroy_distributed_environment()
    with contextlib.suppress(AssertionError):
        torch.distributed.destroy_process_group()
    gc.collect()
    if not is_cpu():
        torch.cuda.empty_cache()

llm = LLM(...)
del llm
cleanup()

[quanshr]

Please see #5716. For now, you have to do it manually as follows (based on the code in tests/conftest.py:

import contextlib
import gc

import torch
from vllm.distributed import (destroy_distributed_environment,
                              destroy_model_parallel)
from vllm.utils import is_cpu

def cleanup():
    destroy_model_parallel()
    destroy_distributed_environment()
    with contextlib.suppress(AssertionError):
        torch.distributed.destroy_process_group()
    gc.collect()
    if not is_cpu():
        torch.cuda.empty_cache()

llm = LLM(...)
del llm
cleanup()

It works! Thanks for your help~

[yanyc428]

Hi~ Does this also work with multiple-GPUs?

[DarkLight1337]

It should work in theory.