技术笔记

vLLM 与 DeepSpeed 入门实验 连接到 Agentic RL

从本地小模型服务和双卡微型训练开始,给出可复用的代码、实验变量和记录表,再梳理生成轨迹、训练更新与权重同步的系统边界。

作者与贡献:基于官方文档与论文整理的原理笔记;公式为理论估算,配套 GPU 实测见独立实验记录。

实验的第一步可以很小:让一个模型正常回答,让一个训练循环正常更新,然后再比较配置。这样遇到错误时,容易判断是环境、接口还是调优引起的。

本篇保留入门方案与示例。已经执行的本机实验、固定版本配置和原始数据请看 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

bash
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,限制上下文长度,先从一个请求开始。首次使用需要下载模型。

bash
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 限制的是模型接受的上下文总长度,要给输入与输出都留空间。命令运行后,在同一台机器的另一个终端发送请求:

bash
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:

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 接管反向与更新。

python
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)
bash
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 权重传输

我希望接下来沿着这个方向补充项目:先完成可检查的生成与训练小实验,再把“任务—轨迹—奖励—更新”连接起来。第一份完整记录的目标,是说清一次循环的时间与显存花在哪里,而不是先追求很大的模型规模。

END OF NOTE✳
← 返回文章列表