【杂记】vLLM如何指定GPU单卡离线推理

写在前面

仅作个人学习与记录用。主要记录vLLM指定GPU单卡/多卡离线推理的方法。

vLLM官方文档中Environment Variables页面有对指定GPU方法的唯一描述：

# used to control the visible devices in the distributed setting
"CUDA_VISIBLE_DEVICES":
lambda: os.environ.get("CUDA_VISIBLE_DEVICES", None),

在vLLM离线推理（Offline Inference）时，可以通过设置tensor_parallel_size = 1/2/3...，来使用默认的单卡GPU或多卡GPU来推理。但是如果想在指定的单卡/多卡GPU中运行vLLM，那么应该如何以及在哪里设置CUDA_VISIBLE_DEVICES？

一般来说，使用下面三种方法就可以了：

shell指定：

CUDA_VISIBLE_DEVICES=3  python train.py

另一种shell指定（不推荐）：

export CUDA_VISIBLE_DEVICES=3  
python train.py

代码内部指定：

import os
os.environ["CUDA_VISIBLE_DEVICES"]="3"

但是在实际执行代码过程中，可能存在失效的情况。即无论怎么修改可见的GPU编号，最后程序都是按照顺序从第0块开始使用。问题出在哪里呢？

假设一共有四卡，先使用nvidia-smi -L查看可用GPU及序号：

GPU 0: GeForce RTX XXX (UUID: xxx)
GPU 1: GeForce RTX XXX (UUID: xxx)
GPU 2: GeForce RTX XXX (UUID: xxx)
GPU 3: NVIDIA XXX (UUID: xxx)

而在代码中测试，会得到：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
import torch
print(torch.cuda.get_device_name(0))  # 返回GPU名称，设备索引默认从0开始
print(torch.cuda.current_device())  # 返回现在使用的GPU索引

输出：
1
GeForce RTX XXX
0

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import torch
print(torch.cuda.get_device_name(0))  # 返回GPU名称，设备索引默认从0开始
print(torch.cuda.current_device())  # 返回现在使用的GPU索引

输出：
NVIDIA XXX
0

这是因为nvidia-smi命令中的GPU序号与代码中的GPU序号是相反的，nvidia-smi的 GPU序号默认使用PCI_BUS_ID，而py文件代码默认GPU序号遵循FASTEST_FIRST。

那么可以修改上述指定方式如下：

shell指定：

CUDA_VISIBLE_DEVICES=3 export CUDA_DEVICE_ORDER="PCI_BUS_ID" python train.py

另一种shell指定（不推荐）：

export CUDA_VISIBLE_DEVICES=3  
export CUDA_DEVICE_ORDER="PCI_BUS_ID"
python train.py

代码内部指定：

import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "3"

另外需要注意，如果你在离线推理时import了pytorch等包，最好将os.environ["CUDA_VISIBLE_DEVICES"] = "3"移到import torch等代码之前，紧随import os之后，即按照如下的方式：

import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]="3"
import torch
......