第 14 期 | Gemma + vLLM:高性能推理服务部署
🎯 学习目标
- 理解 vLLM 库的核心优势,特别是 PagedAttention 和连续批处理(Continuous Batching)机制,及其如何显著提升 LLM 推理性能。
- 掌握使用 vLLM 部署 Google Gemma 模型推理服务的基本流程,包括环境搭建和模型加载。
- 学会通过 vLLM 提供的 OpenAI-Compatible API 端点与部署的 Gemma 服务进行交互,包括文本补全和聊天补全。
- 了解如何配置 vLLM 服务参数以优化 Gemma 模型的推理效率和资源利用。
📖 核心概念讲解
随着大型语言模型(LLM)在生产环境中的广泛应用,如何高效、低成本地提供推理服务成为了一个关键挑战。传统的 LLM 推理服务在处理高并发请求时,往往面临显存碎片化、GPU 利用率低下等问题。vLLM 正是为了解决这些问题而生,它通过一系列创新技术,显著提升了 LLM 推理的吞吐量和效率。
14.1 vLLM 简介:高性能 LLM 推理的利器
vLLM 是一个用于 LLM 推理和服务的开源库,由伯克利大学 RISE 实验室开发。它旨在最大化 LLM 推理的吞吐量,同时最小化延迟。vLLM 的核心优势在于其创新的 KV Cache 管理和请求调度策略。
14.1.1 传统 LLM 推理的瓶颈
在传统的 LLM 推理中,每个请求都需要维护一个键值缓存(KV Cache),用于存储模型在生成每个 token 时的注意力层的键(Key)和值(Value)。这个 KV Cache 会随着生成序列的增长而不断增大,占据大量的 GPU 显存。
- 显存碎片化 (Memory Fragmentation):如果为每个请求预先分配固定大小的 KV Cache,会导致显存浪费,因为实际的序列长度往往小于预设的最大长度。如果动态分配,又容易造成显存碎片化,难以有效利用。
- 低 GPU 利用率 (Low GPU Utilization):传统的批处理(Batching)模式下,如果批次中的某个请求生成结束,GPU 就会等待整个批次完成才能处理下一个批次,导致 GPU 空闲。
- 高延迟 (High Latency):批处理大小的限制和等待机制,都会增加请求的端到端延迟。
14.1.2 vLLM 的核心技术
vLLM 通过引入以下两项关键技术,彻底改变了 LLM 推理的效率:
PagedAttention: PagedAttention 是 vLLM 的核心创新,其灵感来源于操作系统中的虚拟内存分页和内存管理。它将 KV Cache 分割成固定大小的“块”(blocks),并以页表(page table)的形式管理这些块。
- 块的动态分配:与操作系统管理物理内存页类似,PagedAttention 允许将 KV Cache 的块非连续地存储在显存中。当序列增长时,只需分配新的块并更新页表,而无需重新分配整个连续内存区域。
- 显存零碎片:由于块大小固定且可以非连续存储,显存得以更有效地利用,消除了碎片化问题。
- 共享 KV Cache:在某些场景(如并行采样、Beam Search)中,多个序列可能共享相同的前缀。PagedAttention 允许这些序列共享相同的 KV Cache 块,进一步节省显存。
下图展示了 PagedAttention 的基本思想:
+---------------------------------+ | GPU Global Memory (KV Cache Pool) | +---------------------------------+ | Block 0 | Block 1 | Block 2 | ... | Block N | +---------------------------------+ Request 1: Prompt: "Hello, " KV Cache: [Block A] -> [Block B] -> [Block C] (Logical sequence) (Points to Block 0) (Points to Block 5) (Points to Block 2) (Physical storage) Request 2: Prompt: "How are " KV Cache: [Block X] -> [Block Y] (Logical sequence) (Points to Block 1) (Points to Block 4) (Physical storage) Benefits: - No fragmentation: Blocks are fixed size, efficiently packed. - Dynamic allocation: Allocate blocks only when needed. - Sharing: Multiple requests can share common blocks (e.g., for beam search).连续批处理 (Continuous Batching / Orca): 传统的批处理会等待所有序列生成完成或达到最大长度后才开始处理下一批请求。连续批处理则不同,它允许在 GPU 仍在处理当前批次中的序列时,就将新到达的请求或已完成部分生成的序列(等待下一个 token 生成)添加到当前批次中。
- 最大化 GPU 利用率:只要 GPU 有空闲,就可以立即处理新的 token 生成任务,避免了等待空闲。
- 降低平均延迟:请求无需等待整个批次完成,可以更快地得到响应。
- 高吞吐量:GPU 持续工作,大大提高了每秒处理的 token 数量。
14.2 Gemma 与 vLLM 的结合
Google Gemma 模型家族以其轻量级、高性能和基于 Gemini 技术的优势,在开源 LLM 领域备受关注。将 Gemma 与 vLLM 结合,能够充分发挥两者的长处:
- Gemma 的高效架构:Gemma 模型本身在设计上就注重效率,其较小的模型尺寸(如 2B, 7B)和优化的架构使其在推理时具有较低的计算开销。
- vLLM 的推理优化:vLLM 的 PagedAttention 和连续批处理机制能够进一步榨取 GPU 性能,尤其是在高并发场景下,确保 Gemma 模型以最高效的方式响应请求。
这种组合使得开发者能够以更低的成本、更高的吞吐量部署 Gemma 模型,从而更好地服务于各种 AI 应用。
14.3 OpenAI-Compatible API 端点
vLLM 提供了一个与 OpenAI API 兼容的 RESTful 服务端点。这意味着,任何现有的、为 OpenAI API 设计的客户端代码,都可以通过简单地修改 API 基础 URL 来与 vLLM 部署的 Gemma 服务进行交互。这大大降低了迁移和集成的成本,加速了开发流程。
vLLM 支持以下主要 OpenAI API 端点:
/v1/completions:用于文本补全任务。/v1/chat/completions:用于多轮对话和聊天机器人任务。/v1/models:查询可用的模型。
💻 实战演示
本节将演示如何在 GPU 环境下,使用 vLLM 部署 Gemma-7B-IT 模型,并通过 Python 客户端与服务进行交互。
硬件要求提示:Gemma-7B-IT 模型需要至少 16GB 的 GPU 显存。如果您的 GPU 显存不足,可以尝试使用 google/gemma-2b-it 模型,它通常需要 6-8GB 显存。
14.3.1 环境准备与模型下载
首先,我们需要一个具备 CUDA 支持的 Python 环境。推荐使用 conda 创建一个独立环境。
# 1. 创建并激活 conda 环境
conda create -n gemma-vllm python=3.10 -y
conda activate gemma-vllm
# 2. 安装 vLLM
# vLLM 需要特定版本的 PyTorch 和 CUDA 工具包。
# 请根据你的 CUDA 版本选择合适的安装命令。
# 以下示例适用于 CUDA 12.1 和 PyTorch 2.1。
# 详情请参考 vLLM 官方文档:https://docs.vllm.ai/en/latest/getting_started/installation.html
pip install vllm
# 如果你使用的是较新的 PyTorch 版本,例如 PyTorch 2.2+,vLLM 可能需要手动安装特定 CUDA 版本的轮子。
# 例如,对于 CUDA 12.1:
# pip install torch==2.2.2 torchvision==0.17.2 torchaudio==2.2.2 --index-url https://download.pytorch.org/whl/cu121
# pip install vllm==0.3.3 # 选择一个兼容的 vLLM 版本
# 3. 安装 Hugging Face Transformers (vLLM 会用到)
pip install transformers accelerate
# 4. 登录 Hugging Face CLI (Gemma 模型需要同意协议)
# 你需要在 Hugging Face 网站上接受 Gemma 模型的条款和条件。
# 访问 https://huggingface.co/google/gemma-7b-it 并点击 "Agree and access repository"。
# 然后在终端运行:
huggingface-cli login
# 按照提示输入你的 Hugging Face token。
14.3.2 启动 Gemma vLLM 推理服务
我们将部署 google/gemma-7b-it 模型。在终端中执行以下命令来启动 vLLM 服务:
# 启动 vLLM API 服务器
# --model: 指定要加载的 Hugging Face 模型名称
# --tensor-parallel-size: 用于多 GPU 推理的张量并行大小。对于单 GPU 设置,通常为 1。
# --gpu-memory-utilization: 控制 vLLM 可以使用的 GPU 显存比例。默认是 0.9 (90%)。
# --max-model-len: 模型可以处理的最大序列长度 (prompt + generated tokens)。Gemma 7B 的默认上下文是 8192。
# --port: 服务监听的端口。
# --host: 服务监听的 IP 地址。
python -m vllm.entrypoints.api_server \
--model google/gemma-7b-it \
--tensor-parallel-size 1 \
--gpu-memory-utilization 0.9 \
--max-model-len 8192 \
--port 8000 \
--host 0.0.0.0
服务启动后,你将看到类似以下的输出(首次启动会下载模型,可能需要一些时间):
INFO 03-15 10:30:00 api_server.py:148] args: Namespace(host='0.0.0.0', port=8000, model='google/gemma-7b-it', ...
INFO 03-15 10:30:00 tokenizer_loader.py:102] Loading tokenizer from google/gemma-7b-it...
INFO 03-15 10:30:00 llm_engine.py:108] Initializing an LLM engine with config: ...
INFO 03-15 10:30:00 engine.py:108] # GPU: 1, # CPU: 1
INFO 03-15 10:30:00 engine.py:109] Using Tensor Parallelism with size 1
INFO 03-15 10:30:00 engine.py:110] Loading model weights from google/gemma-7b-it...
... (模型加载和初始化过程) ...
INFO 03-15 10:30:30 engine.py:118] Total KV cache memory: 14.00 GiB
INFO 03-15 10:30:30 engine.py:119] vLLM engine is ready.
INFO 03-15 10:30:30 api_server.py:171] Started server process [PID: 12345]
INFO: Started server process [12345]
INFO: Waiting for application startup.
INFO: Application startup complete.
INFO: Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit)
现在,Gemma 推理服务已经在 http://0.0.0.0:8000 上运行。
14.3.3 使用 Python 客户端进行文本补全 (Completions)
在新开一个终端窗口,激活相同的 conda 环境,并创建一个 Python 脚本 client_completions.py:
# client_completions.py
import requests
import json
import time
# vLLM 服务地址
API_URL = "http://localhost:8000/v1/completions"
HEADERS = {"Content-Type": "application/json"}
MODEL_NAME = "google/gemma-7b-it" # 确保与服务器加载的模型名称一致
def generate_completion(prompt_text: str, max_tokens: int = 100, temperature: float = 0.7, stream: bool = False):
"""
向 vLLM 服务发送文本补全请求。
"""
data = {
"model": MODEL_NAME,
"prompt": prompt_text,
"max_tokens": max_tokens,
"temperature": temperature,
"stream": stream,
"stop": ["<eos>"] # Gemma 的结束标志
}
print(f"\n--- 发送请求 ---")
print(f"Prompt: {prompt_text[:50]}...")
start_time = time.time()
if stream:
print("Response (streaming):")
full_response_content = ""
try:
with requests.post(API_URL, headers=HEADERS, data=json.dumps(data), stream=True) as response:
response.raise_for_status() # 检查HTTP错误
for chunk in response.iter_lines():
if chunk:
decoded_chunk = chunk.decode('utf-8')
if decoded_chunk.startswith("data:"):
try:
json_data = json.loads(decoded_chunk[5:])
if 'choices' in json_data and len(json_data['choices']) > 0:
text_content = json_data['choices'][0].get('text', '')
full_response_content += text_content
print(text_content, end='', flush=True)
except json.JSONDecodeError:
print(f"JSON Decode Error: {decoded_chunk}")
pass
except requests.exceptions.RequestException as e:
print(f"请求发生错误: {e}")
return None
print("\n")
end_time = time.time()
print(f"生成完成。耗时: {end_time - start_time:.2f} 秒")
return full_response_content
else:
try:
response = requests.post(API_URL, headers=HEADERS, data=json.dumps(data))
response.raise_for_status() # 检查HTTP错误
response_json = response.json()
end_time = time.time()
print(f"生成完成。耗时: {end_time - start_time:.2f} 秒")
print("Response (full):")
print(json.dumps(response_json, indent=2, ensure_ascii=False))
return response_json
except requests.exceptions.RequestException as e:
print(f"请求发生错误: {e}")
return None
if __name__ == "__main__":
# 非流式请求
generate_completion(
prompt_text="Write a short, inspiring poem about the beauty of nature and technology.",
max_tokens=150,
stream=False
)
# 流式请求
generate_completion(
prompt_text="Describe the process of photosynthesis in simple terms.",
max_tokens=200,
stream=True
)
print("\n--- 演示并发请求 (简化版) ---")
# 为了演示连续批处理的优势,可以同时运行多个此脚本,或者在同一个脚本中用多线程/asyncio发送多个请求。
# 这里我们模拟一个简单的并发场景,发送两个请求。
# 在实际生产环境中,并发请求会更复杂,需要使用 asyncio 或线程池。
prompts = [
"Explain the concept of quantum entanglement.",
"Write a creative tagline for a new coffee shop called 'Bean Dream'."
]
import threading
def run_request(prompt):
print(f"\n[Thread {threading.get_ident()}] Starting request for: {prompt[:30]}...")
generate_completion(prompt, max_tokens=100, stream=True)
print(f"[Thread {threading.get_ident()}] Finished request for: {prompt[:30]}...")
threads = []
for p in prompts:
thread = threading.Thread(target=run_request, args=(p,))
threads.append(thread)
thread.start()
# 稍微间隔一下,模拟请求错峰到达
time.sleep(0.5)
for thread in threads:
thread.join()
print("\n--- 所有并发请求完成 ---")
运行脚本:
python client_completions.py
你将看到 Gemma 模型生成的文本输出,并且可以观察到流式输出(stream=True)时 token 是一个接一个地打印出来的。在并发请求部分,vLLM 服务器会高效地调度两个请求的生成过程,展示其连续批处理的能力。
14.3.4 使用 Python 客户端进行聊天补全 (Chat Completions)
Gemma-IT 模型是为指令遵循和聊天而训练的。vLLM 也支持 OpenAI 兼容的聊天补全 API。创建一个 Python 脚本 client_chat.py:
# client_chat.py
import requests
import json
import time
# vLLM 服务地址
API_URL = "http://localhost:8000/v1/chat/completions"
HEADERS = {"Content-Type": "application/json"}
MODEL_NAME = "google/gemma-7b-it" # 确保与服务器加载的模型名称一致
def generate_chat_completion(messages: list, max_tokens: int = 150, temperature: float = 0.7, stream: bool = False):
"""
向 vLLM 服务发送聊天补全请求。
"""
data = {
"model": MODEL_NAME,
"messages": messages,
"max_tokens": max_tokens,
"temperature": temperature,
"stream": stream
}
print(f"\n--- 发送聊天请求 ---")
start_time = time.time()
if stream:
print("Response (streaming):")
full_response_content = ""
try:
with requests.post(API_URL, headers=HEADERS, data=json.dumps(data), stream=True) as response:
response.raise_for_status()
for chunk in response.iter_lines():
if chunk:
decoded_chunk = chunk.decode('utf-8')
if decoded_chunk.startswith("data:"):
try:
json_data = json.loads(decoded_chunk[5:])
if 'choices' in json_data and len(json_data['choices']) > 0:
delta = json_data['choices'][0].get('delta', {})
if 'content' in delta:
text_content = delta['content']
full_response_content += text_content
print(text_content, end='', flush=True)
except json.JSONDecodeError:
print(f"JSON Decode Error: {decoded_chunk}")
pass
except requests.exceptions.RequestException as e:
print(f"请求发生错误: {e}")
return None
print("\n")
end_time = time.time()
print(f"生成完成。耗时: {end_time - start_time:.2f} 秒")
return full_response_content
else:
try:
response = requests.post(API_URL, headers=HEADERS, data=json.dumps(data))
response.raise_for_status()
response_json = response.json()
end_time = time.time()
print(f"生成完成。耗时: {end_time - start_time:.2f} 秒")
print("Response (full):")
print(json.dumps(response_json, indent=2, ensure_ascii=False))
return response_json
except requests.exceptions.RequestException as e:
print(f"请求发生错误: {e}")
return None
if __name__ == "__main__":
# 单轮对话
generate_chat_completion(
messages=[
{"role": "user", "content": "Hello, how are you today?"}
],
max_tokens=50,
stream=False
)
# 多轮对话 (流式)
generate_chat_completion(
messages=[
{"role": "user", "content": "Tell me a fun fact about space."},
{"role": "assistant", "content": "Did you know that there are more stars in the universe than grains of sand on all the beaches on Earth? That's over a sextillion stars!"},
{"role": "user", "content": "Wow, that's amazing! Can you tell me another one?"}
],
max_tokens=100,
stream=True
)
# 包含系统消息的对话 (Gemma 通常不直接支持系统消息,但 vLLM 接口兼容)
# 对于 Gemma,系统消息通常会被合并到第一个用户消息中,或者在提示词模板中处理。
generate_chat_completion(
messages=[
{"role": "system", "content": "You are a helpful assistant that provides concise answers."},
{"role": "user", "content": "What is the capital of France?"}
],
max_tokens=30,
stream=False
)
运行脚本:
python client_chat.py
通过这些实战演示,你应该能够成功部署 Gemma 推理服务,并使用 OpenAI 兼容的 API 进行文本和聊天补全。vLLM 在后台默默工作,通过 PagedAttention 和连续批处理确保了高效率的推理。
🔧 涉及的工具与命令
| 工具/命令 | 描述 |
|---|---|
conda |
Python 环境管理工具,用于创建和管理独立的 Python 虚拟环境。 |
pip install vllm |
安装 vLLM 库及其依赖。 |
huggingface-cli login |
Hugging Face 命令行工具,用于登录 Hugging Face 账户以下载私有或受限模型(如 Gemma)。 |
python -m vllm.entrypoints.api_server |
启动 vLLM RESTful API 服务器。 |
--model |
vLLM 服务器参数,指定要加载的 Hugging Face 模型名称。 |
--tensor-parallel-size |
vLLM 服务器参数,配置张量并行大小,用于多 GPU 推理。 |
--gpu-memory-utilization |
vLLM 服务器参数,设置 GPU 显存使用比例。 |
--max-model-len |
vLLM 服务器参数,设置模型最大上下文长度。 |
--port |
vLLM 服务器参数,指定服务监听的端口。 |
--host |
vLLM 服务器参数,指定服务监听的 IP 地址。 |
requests |
Python HTTP 库,用于向 vLLM API 服务发送请求。 |
json |
Python 内置库,用于处理 JSON 格式数据。 |
📝 本期要点回顾
- vLLM 核心优势:vLLM 通过 PagedAttention 和连续批处理两大创新技术,显著提升了 LLM 推理服务的吞吐量和效率,解决了传统推理中显存碎片化和 GPU 利用率低下的问题。
- PagedAttention:类似于操作系统的虚拟内存分页机制,将 KV Cache 划分为固定大小的块,实现显存的零碎片化和动态按需分配,并支持 KV Cache 共享。
- 连续批处理 (Continuous Batching):允许 GPU 持续处理请求,无需等待整个批次完成,从而最大化 GPU 利用率并降低平均延迟。
- Gemma + vLLM 协同:将 Google Gemma 模型与 vLLM 结合,能够充分发挥 Gemma 的高效架构和 vLLM 的推理优化能力,实现高性能、低成本的 LLM 推理部署。
- OpenAI-Compatible API:vLLM 提供与 OpenAI API 兼容的 RESTful 端点,极大地简化了现有 LLM 应用的集成和迁移工作。
- 部署实践:通过
python -m vllm.entrypoints.api_server命令可以轻松启动 Gemma 推理服务,并通过requests库发送文本补全和聊天补全请求。
🔗 参考资料
- vLLM 官方 GitHub: https://github.com/vllm-project/vllm
- vLLM 官方文档: https://docs.vllm.ai/en/latest/
- Google Gemma 官方 GitHub: https://github.com/google-deepmind/gemma
- Google AI Gemma 生态页面: https://ai.google.dev/gemma