如何用一条脚本跑完 ChatGPT 类模型的 RLHF 训练：DeepSpeed-Chat 与 Hybrid Engine 实战指南

如何用一条脚本跑完 ChatGPT 类模型的 RLHF 训练DeepSpeed-Chat 与 Hybrid Engine 实战指南【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed9 小时、Azure 云上 300 美元以内一个 13B 参数的 ChatGPT 风格模型就训完了——这是 DeepSpeed-Chat 给出的实测数字。它是构建在 Hybrid Engine 之上的 RLHF 训练流水线把 DeepSpeed 的训练与推理能力装进同一个引擎让同一个模型在生成经验和更新权重两种模式之间无缝切换。下面按流水线长什么样 → 怎么跑通 → 快在哪 → 选多少卡划算的顺序把它拆开讲一遍。先看流水线从预训练到 ChatGPT 的三步DeepSpeed-Chat 复刻了 InstructGPT 的三步走。Step 1 拿人工精选的问答对微调预训练 LLMSFTStep 2 训一个更小的奖励模型 RW数据是人类对同一问题多个回答的排序Step 3 用 PPO 把 RW 的打分反馈灌回去继续训 SFT 模型。这里有两个容易忽略、却常被其他实现省掉的设计一个是可选的 EMA checkpoint最终模型取 αEMA(1−α)Actor 这类加权平均回答质量通常好于裸的最终模型另一个是混合训练把预训练的下一词预测目标和 PPO 目标掺在一起防止模型在 SQuAD2.0 这类公开基准上掉点。官方还配了数据抽象与混合层多份数据集可以先统一格式再切分到三个阶段。一条 train.py 跑通全流程示例项目DeepSpeedExamples 里的 DeepSpeed-Chat 应用目录把三阶段打包进一个 train.py换 actor 模型和部署形态两个参数就能覆盖不同规模pip install deepspeed0.9.0 pip install -r requirements.txt # 8 卡单节点13B actor 350M 奖励模型 python train.py --actor-model facebook/opt-13b --reward-model facebook/opt-350m --deployment-type single_node # 64 卡多节点8 个 DGX 节点 python train.py --actor-model facebook/opt-66b --reward-model facebook/opt-350m --deployment-type multi_node # 单卡消费级 GPU 试跑 1.3B python train.py --actor-model facebook/opt-1.3b --reward-model facebook/opt-350m --deployment-type single_gpu这段只做一件事安装依赖后按部署形态启动三阶段训练。耗时参考Step 3 是大头OPT-13B 在 8×A100-40G 上约 2.5h 0.25h 10.8h合计 13.6hOPT-66B 三步 82 分钟 5 分钟 7.5 小时总共约 9hOPT-1.3B 在单张 A6000-48G 上约 2900s 670s 1.2h一顿午餐的功夫就有可玩的 checkpoint。训完可直接用内置推理 API 体验多轮对话。双模式怎么切换Hybrid Engine 源码拆解一个引擎为什么能同时伺候训练和生成答案在 deepspeed/runtime/hybrid_engine.py 的DeepSpeedHybridEngine它直接继承标准DeepSpeedEngine初始化时按推理策略把 Transformer 层、nn.Linear、nn.Embedding、nn.LayerNorm等替换成推理容器同时用_orig_modules和_orig_fwds存下原始模块与前向函数——随时切回训练路径靠的就是这份备份。每个迭代的关键动作在generate()里。进入生成前它先记录全局 batch 规模、启动计时把 gather 与纯生成时延分开统计若开了release_inference_cache就用retake_inference_cache()重新申请推理 workspace生成结束后调workspace.release_workspace()把显存还给训练阶段。如果模型类型匹配不到任何推理策略还有兼容兜底打印警告并回退到模型原生generate()。开启enable_cuda_graph时会在 decode 侧构建DecodeGraphCache复用图缓存见 deepspeed/runtime/hybrid_engine_graph.py初始化时自动校验 ZeRO 阶段是否支持。API 层面一个迭代就是两个调用——这正是推理/训练双阶段抽象在接口上的落点for prompt_batch in prompt_train_dataloader: out trainer.generate_experience(prompt_batch) # 推理路径生成经验 actor_loss, critic_loss trainer.train_rlhf(out) # 训练路径PPO 更新 actor 与 critic显存优化ZeRO、LoRA 与张量并行怎么叠PPO 阶段的显存压力来自 SFT 模型与奖励模型的多个副本全程驻留Hybrid Engine 的做法是叠三层互不冲突的优化。ZeRO LoRAgenerate()的分支里非 ZeRO-3 场景下生成前fuse_lora_weight()把 LoRA 权重融合进推理容器、生成后unfuse_lora_weight()还原——推理用融合权重训练用分片状态不多占一份。ZeRO-3 pin_parametersgather_all_layers由pin_parameters决定生成前用GatheredParameters把参数拉全当inference_tp_size 1时不一次性全 gather而是按tp_gather_partition_size默认 8 层一组分组 gather逐组应用张量并行把峰值显存压下来。还有一个值得注意的取舍生成阶段横向扩展用的是张量并行而不是 ZeRO。生成是内存带宽受限的TP 减少卡间通信、保住带宽利用率ZeRO 分片留给训练阶段。两种切分方式在同一个引擎里自动切换。Hybrid Engine 配置项逐条讲字段全部定义在 deepspeed/runtime/config.py 的HybridEngineConfig写进 DeepSpeed JSON 的hybrid_engine块即可配置字段默认值作用与使用场景enabledFalseHybrid Engine 总开关max_out_tokens512生成最大长度同时作为推理容器的 min/max 输出长度inference_tp_size1推理张量并行规模1 走 TP 分组 gatherrelease_inference_cacheFalse生成后释放 workspace、训练前再申请压显存峰值pin_parametersTrueZ3 下生成前 gather 全部非 TP 参数并驻留tp_gather_partition_size8Z3TP 时按每 8 层分组 gather 的步长enable_cuda_graphFalsedecode 阶段 CUDA Graph 缓存初始化时校验 ZeRO 阶段仓库里还有一份可跑的最小样例 tests/hybrid_engine/hybrid_engine_config.jsontrain_batch_size 32、micro 2、stage 0 配 CPU 参数 offload、fp16、梯度裁剪 1.0配套的 tests/hybrid_engine/hybrid_engine_test.py 验证了训练—生成切换下的端到端可用性。快在哪端到端吞吐对比 快的来源很直接一次 RLHF 迭代里大部分时间耗在生成上以 1.3B 单次迭代的耗时分解为例生成占大头。换掉 DeepSpeed 的高性能推理内核后生成阶段相对 HuggingFace 最高 9 倍吞吐、相对 Colossal-AI 达 15 倍端到端看单卡 A100-40G 相对其他 RLHF 系统提升超 10 倍8 卡单节点相对 Colossal-AI 有 6–19 倍、相对 HF DDP 有 1.4–10.5 倍加速图中无柱形即 OOM。规模差距更大同一硬件上 Colossal-AI 单卡最大 1.3B、单节点 A100-40G 最大 6.7B而 DeepSpeed-HE 分别到 6.5B 和 50B最大规模扩了 7.5 倍。实测其有效吞吐比现有系统高约 19 倍——换句话说别的系统当时只跑在峰值的 5% 以下。成本与选卡数的拐点先上成本数据均为 Step 3基准规格见最后一节GPU 配置6.7B13B30B66B175B8×A100-40G5.7h10.8h1.85 天——8×A100-80G4.1h / $1329h / $29018h / $5802.1 天 / $1620—64×A100-80G—1.25h / $3204h / $10247.5h / $192020h / $5120单卡能扛多大靠 Hybrid Engine 的显存管理A100-80G 单卡可训 13BV100-32G 也能跑 2.7B。再来看该用几张卡。生成阶段约占 20% 计算量、训练占 80%但生成是带宽受限的实际会占掉大部分墙钟时间。扩展曲线13B/66B actor 350M 奖励模型随 DGX 节点数增加有一个明显拐点小规模下卡越多、单卡显存越宽裕单卡 batch 可以开更大呈现超线性扩展大规模下全局 batch 上限示例场景 1024 组、序列长 512封顶了单卡 batch退化成近线性甚至次线性。所以最佳点就在两段的分界处——给定最大全局 batch选batch 刚好放不下更大的卡数成本效率最高。适用前提引用这些数字前先看这里⚠️ 全文成本与耗时都基于官方 2023 年基准规格总计 135M tokens 训一个 epoch其中 67.5M query tokens131.9k 条、长 256与 67.5M 生成 tokens131.9k 条回答、长 256每步最大全局 batch 为 0.5M tokens即 1024 组 query-answer 对。你的数据规模或序列长度不同时直接套用这些数字会偏差明显。另外仓库当前的 Hybrid Engine 已在博客版本上演进CUDA Graph、共享 prefill workspace 等能力具体字段与行为以源码和测试为准引用论文请用 arXiv:2308.01320。【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考