DeepSpeed-Chat：单卡训 13B RLHF 模型，1 条命令跑完三阶段

DeepSpeed-Chat：单卡训 13B RLHF 模型，1 条命令跑完三阶段

【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed

DeepSpeed-Chat 让 RLHF（基于人类反馈的强化学习）三阶段训练可以一条命令跑完，单张 A100-80G 就能训 13B 模型。读完你能带走三样东西：一条可直接复制的 train.py 命令、hybrid_engine 配置块的全部 7 个字段及默认值、一组实测数字（8×A100-80G 上 9 小时训完 OPT-13B，Azure 成本约 $290，来源：官方博客 Table 1）。核心引擎叫 Hybrid Engine：同一个模型在"训练"和"生成"两条执行路径间切换。

三层结构：数据、训练，和一个会换路的引擎

整个系统分三层。数据层统一不同数据源的格式，再切分混合进三个阶段；训练层按 InstructGPT 流程走：Step 1 监督微调（SFT）、Step 2 奖励模型微调、Step 3 用 PPO（一种限制策略更新幅度的强化学习算法）微调 actor 模型，另附 EMA checkpoint 与混合训练两个可选特性。前两步是普通微调，由 ZeRO（把优化器状态切到多卡以省显存）驱动；第三步每次迭代要做两件事——先生成回答、再训练 actor 和奖励模型。Hybrid Engine 专门接管第三步：训练时它就是普通 DeepSpeed 引擎，生成时切换成推理内核，像同一台发动机在"烧油"和"发电"两种模式间切换，不需要重建模型或搬运权重。

最小路径：2 小时跑通 1.3B 示例

第一步，克隆本仓库并装依赖，约 5 分钟，任意 CPU 机器即可：

git clone https://gitcode.com/GitHub_Trending/de/DeepSpeed.git cd DeepSpeed pip install -r requirements.txt python setup.py install

装完应看到deepspeed --version有正常版本号输出。第二步，先跑仓库自带的 Hybrid Engine 测试验证训练—生成切换可用，单卡 10 分钟内结束，跑完应看到1 passed：

python -m pytest tests/hybrid_engine/hybrid_engine_test.py

第三步，正式训练。train.py 示例脚本位于官方 DeepSpeedExamples 仓库的 DeepSpeed-Chat 应用目录，克隆后装依赖，单卡 48G 显存约 2.2 小时出 checkpoint（来源：官方博客 Table 6）：

python train.py --actor-model facebook/opt-1.3b \ --reward-model facebook/opt-350m --deployment-type single_gpu

跑完后用推理 API 交互，多轮对话输出长这样：

Human: Do you know Microsoft? Assistant: Yes, I'm a big fan of Microsoft. What about them? Human: Can you explain it to a 6-year old child? Assistant: Sure... Microsoft is a company that makes computers, and they make a program called "Windows"...

不同规模与资源的需求（来源：官方博客 Table 1/4/6）：

模型	硬件	三阶段总耗时
OPT-1.3B	单张 A6000 48G	约 2.2 小时
OPT-6.7B	8×A100-80G	4.1 小时（$132）
OPT-13B	8×A100-80G	13.6 小时（Step 3 占 10.8 小时）
OPT-66B	64×A100-80G	约 9 小时

想自定义训练策略时，示例仓库给出的 API 入口只有两个动作：

engine = DeepSpeedRLHFEngine(actor_model_name_or_path=path, tokenizer=tokenizer, args=args) trainer = DeepSpeedPPOTrainer(engine=engine, args=args) for prompt_batch in prompt_dataloader: out = trainer.generate_experience(prompt_batch) # 生成阶段：推理加速路径 actor_loss, critic_loss = trainer.train_rlhf(out) # 训练阶段：PPO 更新

（该 API 在 DeepSpeedExamples 的 DeepSpeed-Chat 应用中，本仓库提供其背后的引擎实现。）

Hybrid Engine 里值得看的四处实现

generate()：同一模型上训练与生成的切换开关

它做什么：同一个引擎实例在不重建模型、不复制权重的前提下切换执行路径。

怎么做到的：初始化时预建推理容器并保存原始 forward；generate()按 ZeRO 阶段把分片参数 gather 进显存、融合 LoRA、走推理内核，结束后解融合并回收推理 workspace。

源码落点：deepspeed/runtime/hybrid_engine.py（DeepSpeedHybridEngine.generate，类定义在 L31-L75）。这段代码的关键行为是：L243-317 是 ZeRO-3 分支（GatheredParameters逐层 gather 参数后推理），L318-329 是非 ZeRO-3 分支（直接 fuse LoRA → 生成 → unfuse），L332-335 在release_inference_cache开启时调用workspace.release_workspace()把显存还回训练阶段。

一个具体例子：配置为 ZeRO-3 且inference_tp_size=2时，生成前引擎按每 8 层一组（tp_gather_partition_size默认 8）gather 非 TP 参数；生成后对每层推理容器调用release_memory()（L297-298）释放 KV-Cache 与中间结果占用的显存。

容易踩的坑：LoRA 的融合与还原必须配对。ZeRO-3 下还原走的是unfuse_lora_weight_non_pinned()（L170-177），只对已 gather 的参数逐层 unfuse，若误用非 pin 路径会残留未还原的参数。

hybrid_engine配置块的全部字段定义在 deepspeed/runtime/config.py，可直接写进 DeepSpeed JSON 配置，仓库样例见 tests/hybrid_engine/hybrid_engine_config.json：

字段	类型/默认值	实际作用
enabled	bool /False	开启 Hybrid Engine
max_out_tokens	int /512	生成最大长度，同时作为推理容器的 workspace 上限
inference_tp_size	int /1	推理阶段张量并行规模，>1时按此规模分组
release_inference_cache	bool /False	生成后释放 workspace，训练前重新申请
pin_parameters	bool /True	ZeRO-3 下生成前 gather 全部非 TP 层参数
tp_gather_partition_size	int /8	ZeRO-3 + TP 场景按每 8 层分组 gather 的步长
enable_cuda_graph	bool /False	启用 decode 阶段 CUDA Graph 缓存

推理容器与张量并行：生成阶段为什么快

它做什么：把训练层的 forward 换成带 KV-Cache 的推理内核，让内存带宽受限的生成阶段跑得更快。

怎么做到的：create_inference_containers()递归遍历模型子模块，命中策略的层（nn.Linear、nn.Embedding、nn.LayerNorm、OPT 位置编码等）替换为推理容器并记录原始 forward；inference_tp_size > 1时，create_inference_module()把全部 rank 按 tp_size 张卡一组建 TP 组，每组持有完整模型，batch 在各组间用 all_gather 分配。

源码落点：deepspeed/runtime/hybrid_engine.py。关键行为是 L386-394 用global_rank // inference_tp_size算出组号、dist.new_group(ranks)建组；生成时 L279-289 把输入张量扩inference_tp_size倍并跨组 all_gather，即每个 TP 组只处理全局 batch 的一部分。

一个具体例子：8 卡 +inference_tp_size=2得到 4 组双卡，每组承担 1/4 的 batch。此时生成吞吐取决于组内 TP 通信而非 ZeRO 的梯度 allreduce——博客实测生成阶段相对 HuggingFace 最高 9 倍、相对 Colossal-AI 15 倍（来源：官方博客 Figure 5）。

CUDA Graph：把 decode 步骤录下来重放

它做什么：消除每个 decode 步骤重复的 kernel launch 开销。

怎么做到的：enable_cuda_graph开启后，初始化时用DecodeGraphCache包裹原始 forward，并先校验 ZeRO 阶段兼容性；generate()入口调用begin_sequence()绑定本次 decode 步数，图被录制后按步重放。

源码落点：deepspeed/runtime/hybrid_engine.py，缓存本体在 deepspeed/runtime/hybrid_engine_graph.py。关键行为是 L75 用max_out_tokens作为 decode 步数上限建缓存，L240-241 在每次生成前重置序列；校验不通过时 L70-71 打印警告并自动退回无图路径，不会报错中断。

一个具体例子：生成 256 token 的回答时max_out_tokens=512即可让图缓存命中；若把生成长度调到 1024，缓存录制成本上升，且 ZeRO-3 等不支持的阶段会看到 "running without CUDA graphs" 警告。

调优提示：CUDA Graph 只覆盖 decode 阶段；开启后先看日志确认没有触发警告回退，否则白录一遍图。

边界与选型：单卡上限 13B，64 卡到 175B

天花板。单卡支持的最大模型：A100-80G 为 OPT-13B，A100-40G 为 OPT-6.7B，V100 32G 只有 2.7B（来源：官方博客 Table 3）。64×A100-80G 训 175B 需 20 小时（约 $5120），但单卡吞吐回落到 1.3B 模型的 1.2 倍，原因是显存撑不起大 batch（来源：官方博客 Figure 6）。以上数据均基于 2023 年基准：135M tokens、单 epoch、每步全局 batch 0.5M tokens，做成本对比时以此规格为准。

怎么调。三个影响最大的旋钮：

• inference_tp_size：模型 6.7B 及以下用 1；30B 以上设 2 或 4。生成阶段是内存带宽受限的，TP 通信量远小于 ZeRO allreduce，比"加卡分片"更划算。
• release_inference_cache：训练显存紧张时设true。代价是每次生成前重新申请 workspace，deepspeed/runtime/hybrid_engine.py 里有一次失败重试逻辑，申请不到会直接抛RuntimeError。
• max_out_tokens：对齐实际回答长度。它同时决定推理 workspace 和 CUDA Graph 缓存上限，设小会截断输出，设大则浪费显存。

选型判断。当你做 OPT 系模型的 PPO 训练、需要 ZeRO + LoRA 且希望一个引擎同时负责训练与生成时，选 DeepSpeed-Chat 的 Hybrid Engine；当你要频繁更换 RL 算法或模型家族时，考虑 HuggingFace 生态的 RLHF 工具链——本仓库的推理容器策略按有限模型结构匹配（见create_inference_containers的策略表匹配逻辑），非覆盖范围回退到模型原生生成路径；当你需要 66B 以上规模时直接按 64 卡 8 节点规划，单节点跑不了。

下一步建议：先跑通tests/hybrid_engine/hybrid_engine_test.py，然后在自己的迭代里观察generate()记录的_generate_latency与_gather_latency（L337-338），确认你的瓶颈在生成还是参数 gather，再决定动哪个旋钮。

【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed