跳到主要内容

快速开始

本章将带你快速上手 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,
)

参数说明

参数类型默认值说明
nint1为每个提示词生成 n 个独立结果
temperaturefloat1.0采样温度,0 表示确定性输出
top_pfloat1.0核采样阈值,只从累积概率 top_p 的 token 中采样
top_kint-1Top-k 采样,只考虑概率最高的 k 个 token
max_tokensint16最大生成 token 数
presence_penaltyfloat0.0主题重复惩罚,范围 -2.0 到 2.0
frequency_penaltyfloat0.0词语重复惩罚,范围 -2.0 到 2.0
repetition_penaltyfloat1.0重复惩罚系数,1.0 表示无惩罚
stopList[str]None遇到这些字符串停止生成
stop_token_idsList[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 的基础用法:

  1. 离线推理:使用 LLM 类加载模型并生成文本
  2. API 服务:通过 vllm serve 启动 OpenAI 兼容的 HTTP 服务
  3. 流式输出:实时获取生成结果,提升交互体验
  4. 批量处理:利用 vLLM 的自动批处理能力提高效率
  5. 采样参数:通过调整参数控制生成行为

掌握了这些基础知识后,你可以进入核心概念章节,深入了解 vLLM 的工作原理和高级特性。