第一次推理:把 vLLM 跑起来,量出你的 tok/s
RL 的墙钟时间大头在采样。先把推理引擎跑起来,测出这张卡的真实吞吐,后面所有时间估算都以它为基准。
学完这节你能做到
- 在 5090 上起一个 vLLM 服务并完成一次批量生成
- 测出 Qwen3-0.6B 在你的机器上的实际生成吞吐(tok/s)
- 知道 gpu_memory_utilization、max_model_len、CUDA graph 各自影响什么
采样是 RL 的墙钟时间大头
一步 GRPO 长这样:
采样 8 个 prompt × 8 条回答 × 平均 400 token → 25600 token
↓
算 reward
↓
一次前向 + 一次反向
↓
权重同步回推理引擎
在单卡 colocate 下,采样通常占掉 60% 到 80% 的时间。也就是说,你优化训练侧的收益远不如优化采样侧。
所以在写任何训练代码之前,先把推理引擎跑起来,量出这张卡的真实吞吐。这个数字后面到处都要用 —— L4 的时间账计算器第一个输入项就是它。
model.generate() 能用,但没有 PagedAttention、没有连续批处理、没有 CUDA graph。同样的卡,吞吐差一个数量级。
RL 里采样占大头,这一个数量级会直接变成「一次实验 40 分钟」和「一次实验 6 小时」的差别。从第一天就用 vLLM。
拿模型
# 国内建议先设镜像
export HF_ENDPOINT=https://hf-mirror.com
pip install -U "huggingface_hub[cli]"
hf download Qwen/Qwen3-0.6B --local-dir ./models/Qwen3-0.6B
0.6B 大约 1.2 GB,很快。顺手把数据集也拿了,下一节要用:
hf download --repo-type dataset openai/gsm8k --local-dir ./data/gsm8k
两种用法
offline:进程内直接调
写实验脚本时最方便,手写训练循环也用这个。
from vllm import LLM, SamplingParams
llm = LLM(
model="./models/Qwen3-0.6B",
gpu_memory_utilization=0.45, # 只给 vLLM 这么多,其余留给训练
max_model_len=2048,
dtype="bfloat16",
)
params = SamplingParams(n=8, temperature=1.0, top_p=1.0, max_tokens=512)
outputs = llm.generate(["3 + 4 * 2 = ?"], params)
for cand in outputs[0].outputs:
print("---", cand.text.strip()[:120])
server:起 HTTP 服务
多卡时训练和推理分在不同卡上会用到。单卡上一般不用,但调试时方便:
vllm serve ./models/Qwen3-0.6B \
--gpu-memory-utilization 0.9 \
--max-model-len 2048 \
--port 8000
压一次吞吐
这是这节课真正的产出物。固定条件、记录数字。
# bench/bench_rollout.py
import time
from vllm import LLM, SamplingParams
PROMPTS = [f"计算并只输出答案:{i} * 37 + 12 = ?" for i in range(64)]
llm = LLM(model="./models/Qwen3-0.6B", gpu_memory_utilization=0.45, max_model_len=2048)
params = SamplingParams(n=8, temperature=1.0, max_tokens=512)
t0 = time.perf_counter()
outputs = llm.generate(PROMPTS, params)
elapsed = time.perf_counter() - t0
gen_tokens = sum(len(c.token_ids) for o in outputs for c in o.outputs)
seqs = sum(len(o.outputs) for o in outputs)
print(f"序列数 : {seqs}")
print(f"生成 token : {gen_tokens}")
print(f"耗时 : {elapsed:.1f} s")
print(f"吞吐 : {gen_tokens / elapsed:.0f} tok/s")
print(f"平均长度 : {gen_tokens / seqs:.0f} token")
跑完把这四个数记下来:
吞吐 ______ tok/s 平均长度 ______ token
显存占用 ______ GiB 引擎启动耗时 ______ s
因为没有这种东西。同一张 5090 上,同一个 0.6B 模型,改一下 max_model_len、并发数或者关掉 CUDA graph,吞吐能差三倍。
给你一个数字只会让你产生虚假的参照。测你自己的。 后面所有时间估算都基于你这个数。
三个真正影响结果的旋钮
gpu_memory_utilization
vLLM 会预先吃掉这个比例的显存,剩下的留给 KV cache。默认 0.9 —— 单卡 colocate 时这会把训练侧饿死。
纯推理测吞吐 → 0.85 ~ 0.9
训推同卡 → 0.3 ~ 0.5,从 0.3 开始试
max_model_len
决定单条序列的上限,也决定 KV cache 块的规划。设得越大,能容纳的并发越少。
不要图省事写 32768。按你实际需要设:数学题任务 2048 就够,长思维链任务再往上加。
CUDA graph
vLLM 默认会捕获 CUDA graph,把逐层 kernel 启动的开销省掉。启动日志里会看到:
INFO Capturing CUDA graphs: 100%|██████████| 35/35
很多 5090 的教程会让你加 --enforce-eager 来绕过启动阶段的报错。能不加就不加 —— 关掉 CUDA graph 会显著掉吞吐。如果不加就报错,优先升级 vLLM 版本。
- 1.测出纯推理场景的吞吐上限
- 2.确认显存占用与 GPU 利用率
- 3.看显存收紧后吞吐掉多少,为 colocate 做准备
吞吐测量演练。先测纯推理的上限,再测 colocate 下的实际值,比较两者的差距。 goals 看目标,hint 要提示。
记进实验笔记
新建一个 NOTES.md,第一条就写这个:
## 环境基线(2026-08-19)
- 卡:RTX 5090 32GB,驱动 580.65
- torch 2.x + cu128,vLLM x.y
- 模型:Qwen3-0.6B bf16
| 场景 | util | 吞吐 tok/s | 显存 GiB |
| --- | --- | --- | --- |
| 纯推理 | 0.45 | 3852 | 14.5 |
| colocate 预留 | 0.30 | 2961 | 9.8 |
引擎启动耗时约 23 s(每次训练脚本重启都要付一次)
调试期间反复重启脚本,光等引擎就浪费大量时间。建议调试阶段用 Jupyter 或 IPython 保持进程常驻,只重跑训练循环那一段。
单卡训推同卡时,vLLM 的 gpu_memory_utilization 该设多少?
为什么这节课要求你必须自己测吞吐,而不是用一个参考值?
这节课的落点
- 采样占 RL 墙钟时间的 60%~80%,优化采样比优化训练划算
- 用 vLLM 不用
model.generate(),吞吐差一个数量级 gpu_memory_utilization是预占不是按需,colocate 时从 0.3 起步max_model_len设得越大能容纳的并发越少,按需求设不要图省事- CUDA graph 能捕获就别用
--enforce-eager关掉 - 产出物:写进
NOTES.md的吞吐、显存、启动耗时四个数
延伸资料
- ·vLLM 文档 ↗
- ·rlforge 配套脚本
scripts/bench_rollout.py