快速开始
本章将带你快速上手 vLLM,从最简单的模型加载开始,逐步掌握离线推理和在线服务两种使用模式。
第一个 vLLM 程序
让我们从一个最简单的例子开始,使用 vLLM 生成文本。
离线推理
离线推理适合批量处理任务,不需要启动 HTTP 服务。
from vllm import LLM, SamplingParams
# 定义提示词
prompts = [
"人工智能是",
"机器学习的主要应用包括",
"深度学习与传统机器学习的区别是",
]
# 配置采样参数
sampling_params = SamplingParams(
temperature=0.8, # 温度参数,控制随机性
top_p=0.95, # 核采样参数
max_tokens=100 # 最大生成 token 数
)
# 加载模型
# 首次运行会自动下载模型(约 500MB)
llm = LLM(model="facebook/opt-125m")
# 执行推理
outputs = llm.generate(prompts, sampling_params)
# 输出结果
for output in outputs:
prompt = output.prompt
generated_text = output.outputs[0].text
print(f"提示词: {prompt}")
print(f"生成结果: {generated_text}")
print("-" * 50)
运行这个程序,你会看到类似以下的输出:
提示词: 人工智能是
生成结果: 一种模拟人类智能的技术,它可以学习、推理和解决问题...
--------------------------------------------------
提示词: 机器学习的主要应用包括
生成结果: 图像识别、自然语言处理、推荐系统、语音识别等...
--------------------------------------------------
代码解析
让我们逐行理解这段代码:
导入模块
from vllm import LLM, SamplingParams
LLM:vLLM 的核心类,负责模型加载和推理SamplingParams:采样参数配置类,控制生成行为
配置采样参数
sampling_params = SamplingParams(
temperature=0.8,
top_p=0.95,
max_tokens=100
)
这些参数控制文本生成的随机性和长度:
temperature(温度):值越高,生成结果越随机;值越低,结果越确定。通常设置在 0.7-1.0 之间。top_p(核采样):只从累积概率超过 p 的最小 token 集合中采样。0.95 表示考虑前 95% 概率的 token。max_tokens:限制生成文本的最大长度,防止生成过长内容。
加载模型
llm = LLM(model="facebook/opt-125m")
这里使用了一个较小的模型(125M 参数)作为示例。vLLM 会自动从 Hugging Face 下载模型。你也可以使用本地路径:
llm = LLM(model="/path/to/your/model")
执行推理
outputs = llm.generate(prompts, sampling_params)
vLLM 会自动对输入进行批处理,充分利用 GPU 并行计算能力。
启动 API 服务
vLLM 提供了与 OpenAI API 兼容的 HTTP 服务,方便集成到现有应用中。
使用命令行启动
# 基本用法
vllm serve facebook/opt-125m
# 指定端口和主机
vllm serve facebook/opt-125m \
--host 0.0.0.0 \
--port 8000
# 启用更多功能
vllm serve facebook/opt-125m \
--tensor-parallel-size 1 \
--max-model-len 2048 \
--gpu-memory-utilization 0.9
调用 API
服务启动后,可以使用任何 HTTP 客户端调用:
使用 curl
curl http://localhost:8000/v1/completions \
-H "Content-Type: application/json" \
-d '{
"model": "facebook/opt-125m",
"prompt": "人工智能是",
"max_tokens": 100,
"temperature": 0.8
}'
使用 Python
from openai import OpenAI
# 配置客户端
client = OpenAI(
base_url="http://localhost:8000/v1",
api_key="not-needed" # vLLM 不需要 API key
)
# 调用 completions API
response = client.completions.create(
model="facebook/opt-125m",
prompt="人工智能是",
max_tokens=100,
temperature=0.8
)
print(response.choices[0].text)
使用 Chat Completions API
# 对话模式
response = client.chat.completions.create(
model="facebook/opt-125m",
messages=[
{"role": "system", "content": "你是一个 helpful 的 AI 助手。"},
{"role": "user", "content": "什么是机器学习?"}
],
max_tokens=200,
temperature=0.7
)
print(response.choices[0].message.content)
流式输出
对于交互式应用,流式输出可以提供更好的用户体验。
API 流式输出
from openai import OpenAI
client = OpenAI(base_url="http://localhost:8000/v1", api_key="not-needed")
# 设置 stream=True
stream = client.completions.create(
model="facebook/opt-125m",
prompt="人工智能是",
max_tokens=100,
stream=True
)
# 逐 token 接收输出
for chunk in stream:
if chunk.choices[0].text:
print(chunk.choices[0].text, end="", flush=True)
批量推理优化
vLLM 的批处理能力是其核心优势之一。以下是一些优化技巧:
动态批处理
vLLM 会自动将多个请求合并处理,无需手动管理批次:
from vllm import LLM, SamplingParams
import time
llm = LLM(model="facebook/opt-125m")
sampling_params = SamplingParams(temperature=0.8, max_tokens=50)
# 准备大量提示词
prompts = [f"这是第 {i} 个测试提示词," for i in range(100)]
# vLLM 会自动进行批处理
start_time = time.time()
outputs = llm.generate(prompts, sampling_params)
end_time = time.time()
print(f"处理 100 个请求耗时: {end_time - start_time:.2f} 秒")
print(f"平均每个请求: {(end_time - start_time) / 100:.2f} 秒")
异步推理
对于高并发场景,可以使用异步 API:
import asyncio
from vllm import LLM, SamplingParams
from concurrent.futures import ThreadPoolExecutor
llm = LLM(model="facebook/opt-125m")
sampling_params = SamplingParams(temperature=0.8, max_tokens=50)
executor = ThreadPoolExecutor(max_workers=4)
async def generate_async(prompt):
"""异步生成文本"""
loop = asyncio.get_event_loop()
result = await loop.run_in_executor(
executor,
lambda: llm.generate([prompt], sampling_params)
)
return result[0]
async def main():
prompts = [f"问题 {i}: 什么是深度学习?" for i in range(10)]
# 并发执行
tasks = [generate_async(p) for p in prompts]
results = await asyncio.gather(*tasks)
for result in results:
print(result.outputs[0].text[:50] + "...")
if __name__ == "__main__":
asyncio.run(main())
常用采样参数详解
SamplingParams 提供了丰富的参数来控制生成行为:
from vllm import SamplingParams
sampling_params = SamplingParams(
# 基础参数
n=1, # 为每个提示词生成 n 个结果
temperature=0.7, # 温度,控制随机性 (0.0-2.0)
top_p=0.95, # 核采样阈值
max_tokens=100, # 最大生成 token 数
# 高级参数
presence_penalty=0.0, # 存在惩罚,减少重复主题
frequency_penalty=0.0, # 频率惩罚,减少重复词语
repetition_penalty=1.0, # 重复惩罚系数
# 控制生成多样性
top_k=-1, # Top-k 采样,-1 表示禁用
min_p=0.0, # 最小概率阈值
# 停止条件
stop=None, # 停止词列表
stop_token_ids=None, # 停止 token ID 列表
# 其他参数
ignore_eos=False, # 是否忽略结束符
skip_special_tokens=True, # 跳过特殊 token
spaces_between_special_tokens=True,
)
参数说明
| 参数 | 类型 | 默认值 | 说明 |
|---|---|---|---|
n | int | 1 | 为每个提示词生成 n 个独立结果 |
temperature | float | 1.0 | 采样温度,0 表示确定性输出 |
top_p | float | 1.0 | 核采样阈值,只从累积概率 top_p 的 token 中采样 |
top_k | int | -1 | Top-k 采样,只考虑概率最高的 k 个 token |
max_tokens | int | 16 | 最大生成 token 数 |
presence_penalty | float | 0.0 | 主题重复惩罚,范围 -2.0 到 2.0 |
frequency_penalty | float | 0.0 | 词语重复惩罚,范围 -2.0 到 2.0 |
repetition_penalty | float | 1.0 | 重复惩罚系数,1.0 表示无惩罚 |
stop | List[str] | None | 遇到这些字符串停止生成 |
stop_token_ids | List[int] | None | 遇到这些 token ID 停止生成 |
采样策略组合
不同的参数组合适合不同的应用场景:
创意写作
# 高温度 + 高 top_p,生成更多样化的内容
SamplingParams(temperature=1.2, top_p=0.95, max_tokens=500)
代码生成
# 低温度 + 存在惩罚,生成更确定、不重复的代码
SamplingParams(temperature=0.2, top_p=0.95, presence_penalty=0.5)
问答系统
# 极低温度,保证答案的确定性
SamplingParams(temperature=0.1, top_p=0.9, max_tokens=200)
对话系统
# 适中温度 + 频率惩罚,避免重复表达
SamplingParams(temperature=0.7, top_p=0.9, frequency_penalty=0.3)
实际应用示例
文本分类
from vllm import LLM, SamplingParams
llm = LLM(model="facebook/opt-125m")
# 设计分类提示词
def classify_text(text):
prompt = f"""请将以下文本分类为:正面、负面或中性。
文本:{text}
分类:"""
sampling_params = SamplingParams(
temperature=0.1, # 低温度,确定性输出
max_tokens=10,
stop=["\n"] # 遇到换行停止
)
output = llm.generate([prompt], sampling_params)
return output[0].outputs[0].text.strip()
# 测试
texts = [
"这个产品太棒了,强烈推荐!",
"质量很差,完全不值这个价",
"今天天气不错"
]
for text in texts:
result = classify_text(text)
print(f"文本: {text}")
print(f"分类: {result}")
print()
信息抽取
from vllm import LLM, SamplingParams
import json
llm = LLM(model="facebook/opt-125m")
def extract_info(text):
prompt = f"""从以下文本中提取人名、地点和时间信息,以 JSON 格式返回。
文本:{text}
JSON:"""
sampling_params = SamplingParams(
temperature=0.1,
max_tokens=200,
stop=["\n\n"]
)
output = llm.generate([prompt], sampling_params)
result = output[0].outputs[0].text.strip()
try:
return json.loads(result)
except:
return {"error": "解析失败", "raw": result}
# 测试
text = "张三和李四于2024年1月15日在北京开会。"
info = extract_info(text)
print(json.dumps(info, ensure_ascii=False, indent=2))
批量处理文件
from vllm import LLM, SamplingParams
import os
llm = LLM(model="facebook/opt-125m")
sampling_params = SamplingParams(temperature=0.7, max_tokens=200)
def process_files(directory):
"""批量处理目录下的所有文本文件"""
results = []
for filename in os.listdir(directory):
if filename.endswith('.txt'):
filepath = os.path.join(directory, filename)
with open(filepath, 'r', encoding='utf-8') as f:
content = f.read()
# 生成摘要
prompt = f"请总结以下文本的主要内容:\n\n{content[:1000]}\n\n摘要:"
results.append((filename, prompt))
# 批量推理
prompts = [p for _, p in results]
outputs = llm.generate(prompts, sampling_params)
# 输出结果
for (filename, _), output in zip(results, outputs):
summary = output.outputs[0].text
print(f"文件: {filename}")
print(f"摘要: {summary}")
print("-" * 50)
# 使用
# process_files('./documents')
小结
本章介绍了 vLLM 的基础用法:
- 离线推理:使用
LLM类加载模型并生成文本 - API 服务:通过
vllm serve启动 OpenAI 兼容的 HTTP 服务 - 流式输出:实时获取生成结果,提升交互体验
- 批量处理:利用 vLLM 的自动批处理能力提高效率
- 采样参数:通过调整参数控制生成行为
掌握了这些基础知识后,你可以进入核心概念章节,深入了解 vLLM 的工作原理和高级特性。