Skip to content

vLLM 推理框架使用指南

内容整理自学习笔记,仅供面试备考参考;不构成录用、培训或考试承诺。

官网:https://vllm.ai/ 文档:https://vllm.readthedocs.io/en/latest/getting_started/installation.html 源码:https://github.com/vllm-project/vllm

1. 引言

1.1 前言

随着大语言模型(LLM)的不断发展,实际上为这些模型提供服务仍然存在挑战,即使在昂贵的硬件上也可能慢得惊人。

来自加州大学伯克利分校的研究者开源了 vLLM,主要用于快速 LLM 推理和服务。vLLM 的核心是 PagedAttention,这是一种新颖的注意力算法,将操作系统的虚拟内存中分页的经典思想引入到 LLM 服务中。

配备了 PagedAttention 的 vLLM 将 LLM 服务状态重新定义:它比 HuggingFace Transformers 提供高达 24 倍的吞吐量,而无需任何模型架构更改。

1.2 为什么需要 vLLM?

vLLM 是一个开源的 LLM 推理和服务引擎,利用全新的注意力算法「PagedAttention」有效地管理注意力键和值。

对比对象吞吐量提升
vs HuggingFace Transformers高达 24 倍
vs HuggingFace TGI高达 3.5 倍

「小羊驼」Vicuna 在 demo 中用到的就是 FastChat 和 vLLM 的集成。vLLM 最大优势在于提供易用、快速、便宜的 LLM 服务。

1.3 vLLM 具有哪些特点?

特点说明
最先进的服务吞吐量业界领先的推理吞吐量
PagedAttention有效管理注意力的键和值
动态批处理请求支持动态 batch
优化好的 CUDA 内核底层 CUDA 优化
与 HuggingFace 无缝集成支持流行的 HF 模型
高吞吐量解码算法支持并行采样、beam search 等
张量并行支持分布式推理
流输出支持流式输出
兼容 OpenAI API可作为 OpenAI API 的替代

1.4 vLLM 支持哪些 HuggingFace 模型?

模型系列示例
GPT-2gpt2、gpt2-xl 等
GPTNeoXEleutherAI/gpt-neox-20b、databricks/dolly-v2-12b、stabilityai/stablelm-tuned-alpha-7b 等
LLaMAlmsys/vicuna-13b-v1.3、young-geng/koala、openlm-research/open_llama_13b 等
OPTfacebook/opt-66b、facebook/opt-iml-max-30b 等

2. vLLM 性能如何?

2.1 实验设置

设置模型硬件
设置一LLaMA-7BNVIDIA A10G GPU
设置二LLaMA-13BNVIDIA A100 GPU (40GB)

2.2 性能对比结果

场景vs HFvs TGI
并行采样高 14-24 倍高 2.2-2.5 倍
Beam Search高 8.5-15 倍高 3.3-3.5 倍

数据来源:从 ShareGPT 数据集中采样输入/输出长度进行测试。

3. vLLM 依赖包

依赖要求
OSLinux
Python3.8 or higher
CUDA11.0 – 11.8
GPUcompute capability 7.0 or higher(V100, T4, RTX20xx, A100, L4 等)

4. vLLM 安装

4.1 构建环境

bash
$ conda create -n py310_chat python=3.10       # 创建新环境
$ source activate py310_chat                   # 激活环境

4.2 使用 pip 安装

bash
$ pip install vllm

4.3 使用 source 安装

bash
$ git clone https://github.com/vllm-project/vllm.git
$ cd vllm
$ pip install -e .  # This may take 5-10 minutes.

5. vLLM 使用

5.1 离线推理

python
from vllm import LLM, SamplingParams

# 定义输入 prompt
prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]

# 采样温度 0.8;核采样(Nucleus Sampling / top_p)设为 0.95
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

# 初始化 vLLM engine
llm = LLM(model="facebook/opt-125m")

# 使用 llm.generate 生成结果
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}")

目前 LLMs 并没有支持所有模型,具体可查看 supported-models。

5.2 在线推理

vLLM 可以作为 LLM 服务进行部署,提供了 FastAPI 服务器示例,位于 vllm/entrypoints/api_server.py。服务器使用 AsyncLLMEngine 类来支持异步处理传入请求。

启动服务:

bash
$ python -m vllm.entrypoints.api_server --model facebook/opt-125m

调用服务:

bash
$ curl http://localhost:8000/generate \
  -d '{
    "prompt": "San Francisco is a",
    "use_beam_search": true,
    "n": 4,
    "temperature": 0
  }'

5.3 OpenAI 兼容服务器

vLLM 可以部署为模仿 OpenAI API 协议的服务器,允许 vLLM 被用作使用 OpenAI API 的应用程序的插入式替换。

启动服务:

bash
$ python -m vllm.entrypoints.openai.api_server --model lmsys/vicuna-7b-v1.3

查询模型列表:

bash
$ curl http://localhost:8000/v1/models

查询补全:

bash
$ curl http://localhost:8000/v1/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "facebook/opt-125m",
    "prompt": "San Francisco is a",
    "max_tokens": 7,
    "temperature": 0
  }'

使用 OpenAI Python 包调用:

python
import openai

# 修改 OpenAI 的 API key 和 API base 以使用 vLLM 的 API 服务器
openai.api_key = "EMPTY"
openai.api_base = "http://localhost:8000/v1"

completion = openai.Completion.create(
    model="facebook/opt-125m",
    prompt="San Francisco is a"
)
print("Completion result:", completion)

6. vLLM 分布式推理与服务

vLLM 支持分布式张量并行推理和服务,目前支持 Megatron-LM 的 tensor parallel algorithm,使用 Ray 管理分布式运行时。

6.1 安装 Ray

bash
$ pip install ray

6.2 多 GPU 推理

python
from vllm import LLM

llm = LLM("facebook/opt-13b", tensor_parallel_size=4)
output = llm.generate("San Franciso is a")

6.3 多 GPU 服务

bash
$ python -m vllm.entrypoints.api_server \
  --model facebook/opt-13b \
  --tensor-parallel-size 4

6.4 跨机器分布式部署

bash
# 在 head 节点上
$ ray start --head

# 在 worker 节点上
$ ray start --address=<ray-head-address>

之后在 head 节点上启动 vLLM 进程,将 tensor_parallel_size 设置为所有机器上的 GPU 总数即可实现多机分布式推理。