vLLM环境安装与运行实例【v0.5.4】

vLLM

vLLM is a fast and easy-to-use library for LLM inference and serving.
vLLM是一个快速易用的LLM推理和服务库。
（注意：vllm使用v0.5.4版本，否则使用PyTorch报错）

vLLM环境安装

1. 基础环境安装

环境	版本	备注
CUDA	12.1	PyTorch==2.4.0需要对应cuda==12.1
Python	3.10.4	可以通过Annaconda或者源码安装
vllm	0.5.4	使用v0.5.4版本，否则使用PyTorch报错
torch	2.4.0	PyTorch==2.4.0需要对应cuda==12.1

具体安装可参考我的另一篇文章fairseq-0.12.2多机训练环境搭建。

2. Python安装（可选，通过Annaconda）

conda create -n vllm python=3.10 -y
conda activate vllm

vllm为conda创建的环境名称

3. 安装vllm

#使用v0.5.4版本，否则使用PyTorch报错
pip3 install vllm==0.5.4 -i https://pypi.tuna.tsinghua.edu.cn/simple
#如果PyTorch版本不正确，可按以下安装
pip3 install torch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 --index-url https://download.pytorch.org/whl/cu121 -i https://pypi.tuna.tsinghua.edu.cn/simple

vLLM运行示例

1. 测试代码

from vllm import LLM, SamplingParams

prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

llm = LLM(model="./opt-125m")

outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

2. 测试结果

CUDA_VISIBLE_DEVICES=0 python3 test.py

说明：

• CUDA_VISIBLE_DEVICES为指定运行显卡号

运行结果：