DeepSpeed-Chat 全解析Hybrid Engine 统一训练与推理的 RLHF 流水线【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeedDeepSpeed-Chat 是 DeepSpeed 生态中面向 ChatGPT 类模型的端到端 RLHF 训练方案其核心是把 DeepSpeed 的训练引擎与推理引擎统一为一个 Hybrid Engine让 SFT、奖励模型微调与 PPO 强化学习三个阶段可以用一个脚本跑完并在单卡上训练 13B 级模型。本文以本仓库 blogs/deepspeed-chat/README.md 博客为主体结合 deepspeed/runtime/hybrid_engine.py、deepspeed/runtime/config.pyHybridEngineConfig定义见第 515-522 行与 tests/hybrid_engine/ 测试做源码级讲解回答三个问题如何低成本跑通Hybrid Engine 的差异化机制在哪大规模下的耗时、成本与扩展性如何为什么 RLHF 需要专门的系统而不是普通微调InstructGPT 式 RLHF 与标准预训练、微调的负载特征完全不同。博客指出用现有系统训练一个 6.7B 的对话模型通常需要昂贵的多 GPU 环境且训练效率往往不足硬件能力的 5%即便有多卡集群现有方案也无法支撑数百亿参数级的 RLHF 训练。其技术根源在 Step 3RLHF 阶段每次迭代都要高效处理两个阶段——推理阶段actor 模型生成 token/经验为训练提供输入与训练阶段更新 actor 与奖励模型权重。它引入两大成本显存成本全程需要同时服务 SFT 模型、奖励模型、参考模型等多个副本EMA 采集与混合训练还会带来额外显存占用生成阶段主导耗时actor 模型要对 256 token 的提示逐 token 生成 256 token生成阶段是内存带宽受限的。虽然按计算量看生成只占约 20%、RL 训练占 80%但端到端时间的大头往往花在生成上——这一点在下文的实测耗时分解中可以直接看到。也就是说通用训练框架按前向反向组织资源天然不适配同一模型反复在生成/训练之间切换的负载。DeepSpeed 的解法是把训练与推理的系统能力组合成统一基础设施 Hybrid EngineDeepSpeed-HE训练模式走 ZeRO 与 LoRA 显存优化生成模式套用推理引擎的 KV-Cache 管理、高性能 Transformer 内核与张量并行并在两种模式间无缝切换模型切分方式与显存布局。单脚本跑通三阶段最小命令集与耗时预期训练示例位于 DeepSpeedExamples 仓库的applications/DeepSpeed-Chat目录。以 OPT-13B 为 actor、OPT-350M 为奖励模型单节点部署的最小命令如下约半天完成 13B 全流程训练并产出 checkpoint# 安装 DeepSpeed 与示例依赖 pip install deepspeed0.9.0 git clone https://gitcode.com/GitHub_Trending/de/DeepSpeed cd DeepSpeed # 或 DeepSpeedExamples 中 DeepSpeed-Chat 应用目录 pip install -r requirements.txt # 单节点 8 卡训练 OPT-13B python train.py --actor-model facebook/opt-13b --reward-model facebook/opt-350m --deployment-type single_node同一脚本通过--actor-model与--deployment-type覆盖不同规模与集群# 多节点 64 卡训练 OPT-66B python train.py --actor-model facebook/opt-66b --reward-model facebook/opt-350m --deployment-type multi_node # 单张消费级 GPU 试跑 OPT-1.3B python train.py --actor-model facebook/opt-1.3b --reward-model facebook/opt-350m --deployment-type single_gpu三档配置的端到端耗时分解博客实测模型组合硬件Step 1Step 2Step 3总计Actor OPT-13B Reward OPT-350M8× A100-40G单 DGX 节点2.5hr0.25hr10.8hr13.6hrActor OPT-66B Reward OPT-350M64× A100-80G8 节点82 min5 min7.5hr约 9hrActor OPT-1.3B Reward OPT-350M单张 A6000 48G2900s670s1.2hr约 2.2hr训练完成后可用 DeepSpeed-Chat 的推理 API 直接做多轮对话式评测验证模型具备对话能力而非单轮问答。对于要自建 RLHF 策略的研究者DeepSpeed-Chat 暴露了通用 API 后端。下面是最小调用结构generate_experience走 Hybrid Engine 的推理加速路径train_rlhf用 PPO 目标更新 actor 与 criticengine DeepSpeedRLHFEngine( actor_model_name_or_pathargs.actor_model_name_or_path, critic_model_name_or_pathargs.critic_model_name_or_path, tokenizertokenizer, num_total_itersnum_total_iters, argsargs) trainer DeepSpeedPPOTrainer(engineengine, argsargs) for prompt_batch in prompt_train_dataloader: out trainer.generate_experience(prompt_batch) actor_loss, critic_loss trainer.train_rlhf(out)三阶段流水线对齐 InstructGPT 的训练配方DeepSpeed-Chat 的流水线完整复刻 InstructGPT 的三步训练并内置了两个常被其他工作省略的可选特性Step 1 监督微调SFT用人工精选的 query-回答对微调预训练模型Step 2 奖励模型微调用人工对同一 query 多个答案的排序数据训练一个通常小于 SFT 模型的奖励模型 RWStep 3 RLHF 训练以 PPO 算法用 RW 的反馈继续微调 SFT 模型。EMA指数移动平均checkpoint可选取 EMA checkpoint 作为最终评测模型。InstructGPT 的经验表明其回答质量通常优于常规最终模型Mixture Training混合训练把下一词预测的预训练目标混入 PPO 目标防止模型在 SQuAD2.0 等公开基准上性能退化数据抽象与混合抽象数据集层统一不同数据源格式并支持多数据集混合后再切分到三个阶段便于用自己的数据训练。从上图可以确认 Step 3 的模型拓扑Actor 之外还有冻结的 Reference model、Critic 与 Reward model 副本这正是 Hybrid Engine 必须解决的多副本显存成本问题。Hybrid Engine一个模型两套引擎的无缝切换Hybrid Engine 的设计核心是actor 模型同一份权重在train()/eval()两种模式下分别挂到训练引擎与推理引擎上。从源码结构看这套同模型双引擎由 deepspeed/runtime/hybrid_engine.py 中的DeepSpeedHybridEngine落地第 31 行起它继承标准DeepSpeedEngine推理容器构建create_inference_module/create_inference_containers第 341-425 行初始化时依据inference_policies遍历模型把匹配的 Transformer 层、nn.Linear、nn.Embedding、nn.LayerNorm及 OPT 位置编码层替换为 DeepSpeed 推理容器同时用_orig_modules/_orig_fwds保存原始模块与前向以便切回训练。找不到匹配策略时打印警告并回退到模型原生generate()路径第 138-147 行这是对未适配模型类型的兼容兜底模式切换eval()第 448-490 行把每层orig_module.forward指到推理容器的 forward 并调用transform_for_inference()train()第 492-504 行做反向恢复并调用transform_for_training()。切换发生在用户熟悉的 PyTorch 模式 API 上训练循环无需感知引擎细节权重同步step()第 506-516 行在每次参数更新后调用推理容器的reset_params()保证下一轮生成用的是最新训练权重训练指标拆分eval()中按_gather_latency/_generate_latency/_training_latency打印每迭代的 E2E 时延分解便于定位生成与训练各自的开销占比。生成路径workspace 回收、ZeRO-3 分区 gather 与 LoRA 融合generate()第 232-339 行是 RLHF 每迭代的耗时热点其内部逻辑按配置分三条路径workspace 生命周期管理进入生成前若启用release_inference_cache通过retake_inference_cache()第 179-190 行重新申请 KV-Cache workspace申请失败会gc.collect()empty_cache()后重试仍失败则抛错生成结束后workspace.release_workspace()把显存归还训练阶段——这是在每种模式下重配置显存系统以最大化可用显存的具体实现第 332-335 行ZeRO-3 张量并行inference_tp_size 1第 243-301 行按tp_gather_partition_size默认 8 层一组分层 gather 非驻留参数与 LoRA 参数逐组apply_tensor_parallelism生成输入经all_gather_into_tensor在 TP 组内广播各 rank 只取回自己那份输出ZeRO-3 pin_parameters第 303-317 行生成前一次性 gather 全部层参数驻留显存生成后还原避免逐层 gather 的反复通信。LoRA 的处理在三种场景下各不相同非 ZeRO-3 时生成前fuse_lora_weight()、生成后unfuse_lora_weight()ZeRO-3 非 pinned 参数时用GatheredParameters按层unfuse_lora_weight_non_pinned()第 170-177 行——印证了博客训练用 ZeRO 分片、推理用 TP并在两者间无缝切换模型切分方式的表述。另外仓库当前实现新增了共享 prefill 工作区prepare_shared_prefill第 192-225 行在共享 prompt 前向之前一次性按source_batch_size * repeats分配 workspacerepeat_shared_prefill_cache第 227-230 行复用已算好的 KV-Cache 展开给多个回答分支避免同一 prompt 的重复 prefill。它明确不支持 ZeRO-3、TP、release_inference_cache与 CUDA Graph 组合。CUDA Graph把千次 kernel 启动压成一次重放当enable_cuda_graph开启时引擎构建 deepspeed/runtime/hybrid_engine_graph.py 中的DecodeGraphCache第 39 行起。文件头部的注释解释了设计约束单个 decode 步会发起约千次 kernel 启动且 kernel 从 host 侧计数器读取当前序列长度——捕获会冻结启动参数因此每个 decode 位置各捕获一张图且 replay 不执行 host 代码eager 与 replay 步不能在同一条序列中混用所以是否在整条生成序列上使用图在begin_sequence中一次性决定第 75-105 行生成长度变化时整组图缓存失效重捕。引擎初始化时hybrid_engine.py 第 68-75 行会自动校验 ZeRO 阶段支持性不支持则降级为无 CUDA Graph 运行并打印警告。hybrid_engine 配置项详解deepspeed/runtime/config.py 第 515-522 行定义了hybrid_engine配置块的完整字段与默认值可直接写入 DeepSpeed JSON 配置配置项类型 / 默认值作用结合源码行为enabledbool默认False开启 Hybrid Enginemax_out_tokensint默认512生成最大输出长度作为推理容器的max/min_out_tokens传入inference_tp_sizeint默认1推理阶段张量并行规模1时按 rank 连续切分组建mp_group第 382-405 行ZeRO-3 下启用分区 gather 路径release_inference_cachebool默认False生成后释放推理 workspace、训练前重新申请缓解生成/训练的显存峰值竞争pin_parametersbool默认TrueZeRO-3 下对应gather_all_layers生成前 gather 全部层参数驻留显存换取推理阶段免 gathertp_gather_partition_sizeint默认8ZeRO-3 TP 推理时按每 8 层分组 gather 的步长enable_cuda_graphbool默认False启用 decode 阶段的逐位置 CUDA Graph 缓存仓库提供了可直接参考的样例 tests/hybrid_engine/hybrid_engine_config.jsontrain_batch_size: 32、train_micro_batch_size_per_gpu: 2、zero_optimization.stage: 0含offload_param.device: cpu与stage3_param_persistence_threshold: 0、fp16.enabled: true、gradient_clipping: 1.0。配套的 tests/hybrid_engine/hybrid_engine_test.py 用deepspeed.initialize(modelmodel, argsargs, enable_hybrid_engineTrue)初始化后做eval()→generate→train()切换验证覆盖了训练-生成切换的端到端可用性。性能与成本实测耗时、云成本与横向对比基准前提博客原文强调以下训练耗时与成本数据均针对Step 3基于 DeepSpeed-RLHF 精选数据集与实际测量吞吐135M tokens 训练一个 epoch其中 67.5M query tokens131.9k 条 query长度 256与 67.5M 生成 tokens131.9k 条回答长度 256每步最大全局 batch 为 0.5M tokens1024 组 query-answer 对。做成本或端到端时间对比前务必注意这一口径。单节点 8× A100 训练耗时与 Azure 近似成本硬件OPT-6.7BOPT-13BOPT-30BOPT-66B8× A100-40GB5.7 hr10.8 hr1.85 天NA8× A100-80GB4.1 hr$1329 hr$29018 hr$5802.1 天$1620多节点 64× A100-80GB 训练耗时与成本硬件OPT-13BOPT-30BOPT-66BOPT-175B64× A100-80G1.25 hr$3204 hr$10247.5 hr$192020 hr$5120单卡可训最大模型单卡即可以训练 13B 级以上模型这是 Hybrid Engine 显存管理能力的直接体现V100 32GA6000 48GA100 40GA100 80G最大支持模型OPT-2.7BOPT-6.7BOPT-6.7BOPT-13B与 Colossal-AI、原生 PyTorch 驱动的 HuggingFace 方案在同一硬件上的 Step 3 端到端吞吐对比如下口径单张 A100-40G 与 8× A100-40G 单 DGX 节点无图标表示 OOM加速来源可以从 1.3B 模型的单迭代耗时分解中看到HF-DDP 与 Colossal-AI 的每序列耗时绝大部分消耗在 Generation 阶段DS-Chat 通过推理内核把该阶段压缩到与 RL 训练同量级有效吞吐分析博客给出的各模型规模在效率最优 GPU 数下的生成/训练/有效吞吐为——OPT-1.3B8 卡23.8/62.9/43.4 TFLOPS、OPT-6.7B8 卡32.9/100.7/71.2、OPT-13B8 卡33.0/97.3/67.1、OPT-30B32 卡48.2/103.6/82.0、OPT-66B16 卡29.3/110.4/67.8、OPT-175B64 卡25.4/74.4/52.8单位 TFLOPS/GPU。据此6.7B-66B 区间效率最高扩到 175B 后因显存限制无法支撑大 batch有效吞吐回落但仍比 1.3B 高 1.2 倍。为最大化有效吞吐DeepSpeed-HE 的策略是两阶段都用尽可能大的 batch生成阶段模型装得进单卡时用高性能内核吃满显存带宽装不下时用张量并行而非 ZeRO扩展——TP 减少了 GPU 间通信、维持高带宽利用率。扩展性特征博客的节点扩展曲线13B/66B actor 350M reward随 DGX 节点数增加显示在最多 64 卡上整体扩展良好且呈现小规模超线性、大规模近线性或次线性的拐点。机制在于显存可用性与最大全局 batch 的相互作用ZeRO 使单卡显存占用随 GPU 数下降、单卡 batch 可做大超线性来源而最大全局 batch本例 1024 组、序列长 512最终限制单卡 batch 上限次线性来源。因此对给定的全局 batch最佳吞吐与成本效率出现在超线性/次线性边界处——该拐点由单卡可运行最大 batch size可用显存与全局 batch 的函数决定为用多少卡最划算的选型提供了依据。适用前提与延伸阅读需要说明的是博客中的成本与耗时数据基于 2023 年发布时的 DeepSpeed-RLHF 精选数据集与训练配方135M tokens、单 epoch、每步 0.5M tokens 全局 batch做对比时应以该基准规格为准。本仓库当前的 Hybrid Engine 实现已在此基础上演进如 CUDA Graph、共享 prefill 工作区、未适配模型的generate()回退等能力具体字段与行为请以仓库当前源码与测试为准。若要引用 DeepSpeed-Chat建议使用 arXiv 论文arXiv:2308.01320的 BibTeXarticle{yao2023dschat, title{{DeepSpeed-Chat: Easy, Fast and Affordable RLHF Training of ChatGPT-like Models at All Scales}}, author{Zhewei Yao and Reza Yazdani Aminabadi and Olatunji Ruwase and Samyam Rajbhandari and Xiaoxia Wu and Ammar Ahmad Awan and Jeff Rasley and Minjia Zhang and Conglong Li and Connor Holmes and Zhongzhu Zhou and Michael Wyatt and Molly Smith and Lev Kurilenko and Heyang Qin and Masahiro Tanaka and Shuai Che and Shuaiwen Leon Song and Yuxiong He}, journal{arXiv preprint arXiv:2308.01320}, year{2023} }小结DeepSpeed-Chat 用三件协同的事解决了 RLHF难上手、太贵、扩不动的问题——单脚本三阶段 推理 API 的低门槛体验、完整对齐 InstructGPT含 EMA 与混合训练、多数据源混合的流水线、把训练与推理统一进 Hybrid Engine 的系统设计。可溯源的仓库文件清单混合引擎核心deepspeed/runtime/hybrid_engine.pygenerate()第 232-339 行、LoRA 融合/还原第 156-177 行、ZeRO-3 分区 gather 第 243-301 行、eval/train 切换第 448-504 行、workspace 回收第 179-190 行CUDA Graph 支持deepspeed/runtime/hybrid_engine_graph.pyDecodeGraphCache第 39 行起、begin_sequence第 75-105 行配置定义deepspeed/runtime/config.py 第 515-528 行HybridEngineConfig与get_hybrid_engine_config测试与配置样例tests/hybrid_engine/hybrid_engine_test.py、tests/hybrid_engine/hybrid_engine_config.json官方博客全文blogs/deepspeed-chat/README.md。【免费下载链接】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），仅供参考