How to get GPU memory footprints when using distributed inference?

[HermitSun]

I want to observe the GPU memory footprint of models when performing inferences.

When I perform inferences on a single GPU, torch.cuda.memory_allocated returns a positive number as expected. But when I perform distributed inferences, torch.cuda.memory_allocated returns 0.

Should I use nvidia-smi or some other techniques to get the GPU memory footprint?

Any help would be appreciated.

[WoosukKwon]

Hi @HermitSun, thanks for trying out vLLM and good question. When using multiple GPUs, vLLM creates 1 worker process per GPU. Thus, if you use 2 GPUs, there will be 3 processes in total and the process running your code will not directly use any GPU. To actually get the number, you will need to insert torch.cuda.memory_allocated inside the Worker class or in the model code.

BTW, you can configure the vLLM's GPU memory usage via the gpu_memory_utilization argument in the LLM class.

[HermitSun]

Thank you for your kind reply. After I insert code inside the Worker, it works.

Maybe we can provide some profiling hooks or decorators, if possible.