实验的第一步可以很小:让一个模型正常回答,让一个训练循环正常更新,然后再比较配置。这样遇到错误时,容易判断是环境、接口还是调优引起的。
本篇保留入门方案与示例。已经执行的本机实验、固定版本配置和原始数据请看 RTX 3050 4GB 实验记录,其中单卡实验与这里的双卡示例分别说明。Linux shell 示例应在独立的 GPU 学习环境中运行;模型服务和分布式训练分别使用环境,避免依赖冲突。
系列回顾:学习路线 · vLLM 原理 · ZeRO 显存分析
先记录运行条件
| 必填项 | 记录内容 |
|---|---|
| 硬件 | GPU 型号、数量、显存、卡间互联 |
| 软件 | 系统、驱动、Python、PyTorch、CUDA、vLLM 或 DeepSpeed 版本 |
| 模型和数据 | 模型仓库与 revision、输入样本、随机种子 |
| 工作负载 | 输入输出 token 数、并发、预热与测量轮数 |
| 产物 | 启动命令、完整配置、原始日志、失败信息 |
vLLM 当前入门页使用 Linux 与 Python 3.10 至 3.13,安装应匹配具体加速器环境。开始前按官方指引准备依赖,再记录实际安装版本。vLLM Quickstart
nvidia-smi
python --version
python -m pip show torch vllm
# 在另一套训练环境中执行
python -m pip show torch deepspeed
ds_report
ds_report 用于检查 DeepSpeed 算子与环境状态;安装方式和所需工具链以官方说明为准。DeepSpeed 安装与环境检查
实验一 启动小模型推理服务
沿用官方 Quickstart 的 Qwen/Qwen2.5-1.5B-Instruct,限制上下文长度,先从一个请求开始。首次使用需要下载模型。
vllm serve Qwen/Qwen2.5-1.5B-Instruct \
--host 127.0.0.1 \
--port 8000 \
--served-model-name study-model \
--max-model-len 2048
max-model-len 限制的是模型接受的上下文总长度,要给输入与输出都留空间。命令运行后,在同一台机器的另一个终端发送请求:
curl http://127.0.0.1:8000/v1/chat/completions \
-H 'Content-Type: application/json' \
-d '{
"model": "study-model",
"messages": [{"role": "user", "content": "用三句话解释 KV Cache。"}],
"temperature": 0,
"max_tokens": 128
}'
这一步检查模型能否加载、聊天模板与接口是否可用。单个非流式请求的总耗时不能用来计算 TTFT,也不能证明高并发吞吐。服务只绑定回环地址;远程操作时通过 SSH 隧道访问该机器。
后续压测可以使用 vllm bench serve,先查看当前版本的 --help,固定数据集与长度,再设置并发和请求速率。官方压测命令
| 对照组 | 改变什么 | 保持什么不变 | 观察什么 |
|---|---|---|---|
| A | 并发 1、4、8 | 模型、输入输出长度 | 吞吐、TTFT、错误率 |
| B | 输入 256、1024 token | 并发、输出长度 | Prefill 与 TTFT |
| C | 共同前缀的冷缓存与预热后 | 其余请求结构 | 命中与 Prefill 耗时 |
这些数值是实验档位,不是推荐最优配置;缓存组尤其要记录前缀是否真正命中。
实验二 用微型模型理解 DeepSpeed 训练接口
下面使用两张 CUDA GPU、FP32、小型 MLP 和随机分类数据,只验证引擎、梯度累积和分布式执行。它不是语言模型训练,也不能用来验证 7B 模型的显存结论。
将以下内容保存为 ds_config.json:
{
"train_micro_batch_size_per_gpu": 2,
"gradient_accumulation_steps": 4,
"train_batch_size": 16,
"fp16": {"enabled": false},
"bf16": {"enabled": false},
"zero_optimization": {
"stage": 2,
"reduce_bucket_size": 1000000,
"allgather_bucket_size": 1000000,
"overlap_comm": false
},
"steps_per_print": 4
}
有效全局批量为 $2\times4\times2=16$:每卡微批量 × 梯度累积步数 × 数据并行卡数。配置中的训练批量字段遵循这一关系。DeepSpeed 批量配置
保存以下文件为 train_demo.py。这里显式传入 PyTorch AdamW,避免把入门重点变成优化器扩展编译;DeepSpeed 接管反向与更新。
import argparse
import os
import deepspeed
import torch
from torch import nn
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", "--local-rank", type=int, default=-1)
args = parser.parse_args()
local_rank = int(os.environ.get("LOCAL_RANK", args.local_rank))
if local_rank < 0:
raise RuntimeError("请用 deepspeed launcher 启动")
torch.cuda.set_device(local_rank)
torch.manual_seed(42)
model = nn.Sequential(nn.Linear(128, 256), nn.ReLU(), nn.Linear(256, 4))
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
engine, _, _, _ = deepspeed.initialize(
model=model,
model_parameters=model.parameters(),
optimizer=optimizer,
config="ds_config.json",
)
torch.manual_seed(42 + torch.distributed.get_rank())
engine.train()
for micro_step in range(16):
x = torch.randn(2, 128, device=engine.device)
target = torch.randint(0, 4, (2,), device=engine.device)
loss = nn.functional.cross_entropy(engine(x), target)
engine.backward(loss)
engine.step()
if torch.distributed.get_rank() == 0:
print(f"micro_step={micro_step} loss={loss.item():.4f}", flush=True)
deepspeed --num_gpus=2 train_demo.py
调用引擎 step 的频率按微步执行,梯度累积边界由引擎管理;不要再额外调用一次 optimizer.step()。DeepSpeed 训练接口
随机样本每次都变化,loss 不要求持续下降。这里应检查进程是否正常结束、梯度累积是否按配置发生、日志是否有异常。只有一张卡时,可把全局批量改为 8 并用 --num_gpus=1 验证接口,但不能据此比较多卡分片收益。
把结果整理成自己的实验记录
下面是新实验可以沿用的空白记录模板;本轮已完成的数值、图表与日志见 本机实测记录。
| 配置 | 预热后步耗时 | 峰值显存及口径 | 是否完成 | 原始日志 |
|---|---|---|---|---|
| 基线 | 待测 | 待测 | 待测 | 待补 |
| 只改变一个配置 | 待测 | 待测 | 待测 | 待补 |
测量 GPU 时间前后要同步,区分预热与稳定阶段;多卡训练同时记录各卡情况,并以整体完成一步所需时间比较。小 MLP 只适合检查流程,正式比较 ZeRO 时应另选能体现状态容量压力的模型,保持有效批量等条件一致。
一次 OOM 也值得记录:在哪个阶段失败,修改了哪个变量,是否恢复运行。这样文章会逐渐从概念整理变成有证据的工程实践。
继续连接到 Agentic RL
flowchart TB
A[任务集合与环境] --> B[生成侧 可使用 vLLM]
B --> C[工具调用与环境反馈]
C --> D[轨迹 奖励 与策略版本]
D --> E[训练侧 可选择 DeepSpeed 等后端]
E --> F[导出或传输更新后的权重]
F --> B
图中的框是系统职责,不代表两个库可以直接拼起来训练。接入框架时,需要核对生成时的权重版本、token 与掩码、奖励、截断规则,以及训练后如何同步权重。vLLM 的权重传输接口是进一步阅读的入口。vLLM 权重传输
我希望接下来沿着这个方向补充项目:先完成可检查的生成与训练小实验,再把“任务—轨迹—奖励—更新”连接起来。第一份完整记录的目标,是说清一次循环的时间与显存花在哪里,而不是先追求很大的模型规模。