DeepSpeed-Chat: Easy, Fast and Affordable RLHF Training of ChatGPT-like Models at All Scales

作者： Zhewei Yao, Reza Yazdani Aminabadi, Olatunji Ruwase, Samyam Rajbhandari, Xiaoxia Wu, Ammar Ahmad Awan, Jeff Rasley, Minjia Zhang, Conglong Li, Connor Holmes, Zhongzhu Zhou, Michael Wyatt, Molly Smith, Lev Kurilenko, Heyang Qin, Masahiro Tanaka, Shuai Che, Shuaiwen Leon Song, Yuxiong He
机构： Deepspeed of Microsoft

主要贡献

本文介绍了DeepSpeed-Chat系统，该系统旨在解决当前缺乏易用、高效且经济实惠的端到端RLHF（带有人类反馈的强化学习）训练管道的问题，特别是针对数十亿参数规模的ChatGPT-like模型。核心问题是现有系统在训练中等规模（如6.7B参数）的ChatGPT模型时需要昂贵的多GPU设置，训练效率通常低于机器能力的5%（如图6所示），且无法支持数百亿参数模型的易用、快速和经济训练。这些限制源于缺乏支持InstructGPT RLHF训练管道的鲁棒系统设计，该管道不同于现有深度学习系统设计的标准预训练和微调管道。

研究目标是使RLHF训练民主化，让AI社区更容易访问，通过提供三个关键能力：（1）易用的ChatGPT-like模型训练和推理体验，使用单一脚本从Huggingface预训练模型开始，通过DeepSpeed-RLHF系统完成InstructGPT的所有三个步骤，并生成自定义ChatGPT-like模型，还提供用于测试对话式交互的推理API；（2）DeepSpeed-RLHF管道，该管道复制InstructGPT的训练管道，包括监督微调（SFT）、奖励模型微调和带有人类反馈的强化学习（RLHF），并提供数据抽象和混合能力以支持多个数据源；（3）鲁棒的DeepSpeed-RLHF系统，该系统将训练和推理优化统一到混合引擎（DeepSpeed-HE）中，能够在RLHF中无缝切换推理和训练模式，利用DeepSpeed-Inference的张量并行和高性能Transformer内核进行生成，同时受益于ZeRO和LoRA的内存优化策略，并在整个RLHF管道中优化内存管理和数据移动。

创新点包括DeepSpeed-HE提供的无与伦比的效率和可扩展性，例如在Azure上训练OPT-13B只需9小时、成本低于300美元，OPT-30B只需18小时、成本低于600美元（表1）；在多节点系统上，13B模型可在1.25小时内训练，175B模型可在一天内完成（表2）；单GPU支持超过13亿参数模型（表3）。这些能力使RLHF训练快速、经济且易于访问，即使资源有限的数据科学家也能创建大型模型，从而促进AI领域的创新和发展。

背景知识/关键观察/设计原则

易用训练和推理体验。 DeepSpeed-Chat提供易用的训练和推理体验，例如使用单一脚本训练OPT-13B和OPT-66B模型。如果时间有限，甚至可以在单个消费级GPU上在两小时内训练OPT-1.3B模型。此外，还演示了如何使用DeepSpeed-Chat的RLHF API开发自定义管道。
训练第一个ChatGPT风格模型的简单性。 一个单一脚本完成RLHF训练的所有三个阶段并生成第一个ChatGPT模型。以预训练OPT-13B作为actor模型和OPT-350M作为奖励模型为例，以下单一脚本生成最终13B ChatGPT风格模型：首先安装deepspeed >=0.9.0，然后克隆github仓库，进入目录并安装requirements.txt。运行python train.py --actor-model facebook/opt-13b --reward-model facebook/opt-350m --deployment-type single_node。在大约半天内，13B模型将完全训练完成，其检查点准备就绪。表4显示了每个步骤的训练时间分解。训练后，使用DeepSpeed-Chat的推理API测试最终模型，支持Q&A和对话式交互，例如人类询问关于Microsoft的问题，助手回应并解释给6岁孩子听。

尝试不同模型大小和配置。 用户可以轻松尝试不同模型大小和配置以满足训练时间、资源和质量要求。例如，使用相同脚本训练更大模型如66B，使用64 GPU：在9小时内，66亿参数ChatGPT模型准备好在前端GUI中服务，如表5所示。如果只有1-2小时，可以训练1.3B模型，使用单一数据集，在消费级GPU上完成，并在午餐后获得检查点（表6）。运行python train.py --actor-model facebook/opt-1.3b --reward-model facebook/opt-350m --deployment-type single_gpu。

使用DeepSpeed-Chat的RLHF API自定义管道。 DeepSpeed-Chat允许用户使用灵活API构建自己的RLHF训练管道，这些API可用于重建自定义RLHF训练策略，提供通用接口和后端以支持各种RLHF算法的研究探索。示例代码包括初始化DeepSpeedRLHFEngine（指定actor模型路径、critic模型路径、tokenizer、总迭代数和args），然后创建DeepSpeedPPOTrainer，对于prompt训练数据加载器的每个批次，生成经验并训练RLHF。

方法细节

完整RLHF训练管道。 为了提供无缝训练体验，我们遵循InstructGPT，并在DeepSpeed-Chat中包含完整的端到端训练管道，如图1所示。我们的管道包括三个主要步骤：步骤1是监督微调（SFT），其中仔细选择人类对各种查询的响应来微调预训练语言模型；步骤2是奖励模型微调，其中训练一个单独的（通常小于SFT的）模型（RW），使用具有人类提供的同一查询多个答案排名的数据集。
步骤3的附加特征。 步骤3是RLHF训练，其中使用Proximal Policy Optimization（PPO）算法【[11], Proximal policy optimization algorithms+2017+arXiv preprint arXiv:1707.06347】进一步微调SFT模型，使用来自RW模型的奖励反馈。我们在步骤3中提供两个附加特征以帮助改善模型质量：指数移动平均（EMA）收集，其中可以选择基于EMA的检查点用于最终评估；混合训练，将预训练目标（即下一个词预测）与PPO目标混合，以防止在公共基准如SQuAD2.0【[12], Know what you don’t know: Unanswerable questions for squad+2018+arXiv preprint arXiv:1806.03822】上的性能回归。
附加特征的必要性。 这两个训练特征，EMA和混合训练，通常被其他最近努力省略，因为它们是可选的。然而，根据InstructGPT，EMA检查点通常提供比常规最终训练模型更好的响应质量，混合训练可以帮助模型保留预训练基准解决能力。因此，我们为用户提供它们，以完全获得InstructGPT中描述的训练体验，并争取更高的模型质量。
支持多个数据资源的便利特征。 除了与InstructGPT论文【[7], Training language models to follow instructions with human feedback+2022+Advances in Neural Information Processing Systems】高度一致外，我们还提供便利特征来支持研究者和从业者使用多个数据资源训练自己的RLHF模型：数据抽象和混合能力，DeepSpeed-Chat能够使用多个数据集训练模型以获得更好的模型质量。它配备了（1）抽象数据集层来统一不同数据集的格式；（2）数据分割/混合能力，以便多个数据集被适当混合然后分割到3个训练阶段。为了说明我们训练管道的有效性，我们使用多轮对话演示模型质量，如体验部分所示。

DeepSpeed混合引擎——统一基础设施以驱动和优化RLHF训练。 指令引导RLHF管道的步骤1和步骤2类似于大型模型的常规微调，它们由DeepSpeed训练中的ZeRO-based优化和并行策略的灵活组合驱动，以实现规模和速度。另一方面，管道的步骤3在性能影响方面是最复杂的部分处理。每个迭代需要高效处理两个阶段：a) 用于令牌/经验生成的推理阶段，产生训练的输入；b) 训练阶段以更新actor和奖励模型的权重，以及它们之间的交互和调度。它引入两个主要成本：（1）内存成本，因为需要在阶段3中服务SFT和RW模型的多个副本；（2）主导的生成阶段，如果没有适当加速，将显著减慢整个阶段3。此外，我们在阶段3中添加的两个重要特征，包括指数移动平均（EMA）收集和混合训练，将产生额外的内存和训练成本。
混合引擎的设计。 为了应对这些挑战，我们将DeepSpeed训练和推理的完整系统能力组合成一个统一的基础设施，我们称之为混合引擎。它利用原始DeepSpeed引擎进行快速训练模式，同时轻松应用DeepSpeed推理引擎进行生成/评估模式，为RLHF训练的阶段3提供显著更快的训练系统。如图2所示，DeepSpeed训练和推理引擎之间的过渡是无缝的：通过为actor模型启用典型的eval和train模式，在运行推理和训练管道时，DeepSpeed选择其不同的优化来更快地运行模型并提高整体系统吞吐量。在其推理执行期间，用于RLHF训练的经验生成阶段，DeepSpeed混合引擎使用轻量级内存管理系统来处理KV-cache和中间结果，连同高度优化的推理适应内核和张量并行实现，与现有解决方案相比，实现吞吐量（tokens-per-second）的显著提升。
训练执行中的优化。 在训练执行期间，混合引擎启用内存优化技术，如DeepSpeed的ZeRO系列技术和Low Rank Adaptation（LoRA）【[9], Lora: Low-rank adaptation of large language models+2021+arXiv preprint arXiv:2106.09685】。我们以兼容方式设计和实现了这些系统优化，以便它们可以组合在一起，在统一的混合引擎下提供最高的训练效率。混合引擎可以无缝更改训练和推理中的模型分区，以支持基于张量并行的推理和基于ZeRO的分片机制用于训练。它还可以重新配置内存系统，以在每个模式中最大化内存可用性。这允许通过避免内存分配瓶颈和支持大批量大小来改善性能。混合引擎集成了DeepSpeed训练和推理的系统技术谱系，推动了现代RLHF训练的边界，并为RLHF工作负载提供了无与伦比的规模和系统效率。

DeepSpeed RLHF：通过混合引擎实现无与伦比的规模和效率。 如讨论所述，DeepSpeed-HE是强大推理和训练系统技术的融合，旨在为DeepSpeed-RLHF管道实现优秀的规模和效率，跨越广泛的硬件，使RLHF训练快速、经济且易于AI社区访问。在效率和经济性方面，如表1所示，DeepSpeed-HE可以在Azure云上仅用9小时训练OPT-13B，成本低于300美元，用18小时训练OPT-30B，成本低于600美元。在速度和可扩展性方面，如表2所示，即使是13B模型也可以在1.25小时内训练，大型175B模型可以使用64 GPU集群在一天内训练。在RLHF的访问性和民主化方面，DeepSpeed-HE支持在单个GPU上训练超过13亿参数的模型，如表3所示。
与其他RLHF系统的吞吐量和模型规模可扩展性比较。 与其他RLHF系统如Colossal-AI【[13], Colossal ai+2022+https://github.com/hpcaitech/ColossalAI】或由原生PyTorch【[14], Pytorch: An imperative style, high-performance deep learning library+2019+Advances in neural information processing systems】驱动的HuggingFace【[6], Huggingface’s transformers: State-of-the-art natural language processing+2019+arXiv preprint arXiv:1910.03771】相比，DeepSpeed-RLHF在系统性能和模型可扩展性方面表现出色：在吞吐量方面，DeepSpeed在单个GPU上为RLHF训练启用超过10倍的改进（图3）。在多GPU设置上，它比Colossal-AI快6-19倍，比HuggingFace DDP快1.4-10.5倍（图4）。在模型可扩展性方面，Colossal-AI可以在单个GPU上运行最大1.3B模型，在单个A100 40G节点上运行6.7B模型，而DeepSpeed-HE可以在相同硬件上分别运行6.5B和50B模型，高达7.5倍更大。因此，通过超过一个数量级的更高吞吐量，DeepSpeed-HE解锁了在相同延迟预算下训练显著更大actor模型的能力，或以超过10倍更低成本训练类似大小的模型，与现有RLHF系统如Colossal-AI或HuggingFace DDP相比。这种效率改进源于DeepSpeed-HE利用DeepSpeed推理优化加速RLHF处理中的RLHF生成阶段的能力。图5显示了1.3B参数模型在RLHF训练迭代中的时间分解：大部分时间用于生成阶段。通过利用DeepSpeed的高性能推理内核，DeepSpeed-HE可以在此阶段实现比HuggingFace高9倍的吞吐量改进，比Colossal-AI高15倍，从而实现无与伦比的端到端效率。

有效吞吐量和可扩展性分析（I）有效吞吐量分析。 DeepSpeed-HE在RLHF训练阶段3的有效吞吐量取决于其在生成和RL训练阶段实现的吞吐量。在我们的RLHF管道中，生成阶段约占总计算的20%，而RL训练阶段占剩余80%（详见基准设置https://github.com/microsoft/DeepSpeedExamples/tree/master/applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/BenckmarkSetting.md）。然而，尽管比例小，前者可能占用端到端时间的大部分，因为它需要为每个生成的256令牌运行actor模型一次，初始prompt为256令牌，使其受内存带宽限制且难以实现高吞吐量。相比之下，RL训练阶段是计算绑定的，使用完整512令牌（prompt和生成）的每个样本仅运行参考actor模型几次前向和后向传递，并能实现良好吞吐量。为了最大化有效吞吐量，DeepSpeed-HE优化两个阶段。首先，它使用尽可能大的批量大小来获得两个阶段的更高效率。其次，在生成阶段，当模型适合单个GPU内存时，它利用高性能Transformer内核来最大化GPU内存带宽利用率，当不适合时，利用张量并行（TP）。在生成阶段使用TP而不是ZeRO来适应模型，可以减少GPU间通信并保持高GPU内存带宽利用率。图6显示了DeepSpeed-HE针对从1.3B到175B模型大小的最佳可实现有效吞吐量（以TFlops/GPU计）。它还显示了生成和训练阶段各自实现的吞吐量。DeepSpeed-HE对于6.7B-66B范围的模型最有效。超过此范围到175B，吞吐量下降，因为内存有限，无法支持更大批量大小，但仍比小型1.3B模型高1.2倍效率。这些巨型模型的每GPU吞吐量可以在扩展到更多GPU时进一步改善，提供更多内存用于更大批量大小。此外，我们指出我们的有效性能比现有系统高19倍，如图4所示，这表明它们以低于峰值的5%运行。这展示了优化RLHF工作负载的挑战以及我们系统尽管挑战的有效性。

有效吞吐量和可扩展性分析（II）可扩展性分析。 不同模型大小的最佳有效吞吐量在不同GPU数量下实现。这部分是因为一些更大模型大小需要更多内存来运行。然而，这种行为很大程度上源于DeepSpeed-HE的可扩展性属性，我们接下来讨论。图7显示DeepSpeed-RLHF在最多64 GPU上实现了良好的整体扩展。然而，如果更仔细观察，它显示DeepSpeed-RLHF训练在小规模上实现超线性扩展，随后在大规模上实现近线性或亚线性扩展。这是由于内存可用性和最大全局批量大小之间的交互。由于DeepSpeed-HE由ZeRO-based技术【[15], Zero: Memory optimizations toward training trillion parameter models+2020+SC20: International Conference for High Performance Computing, Networking, Storage and Analysis】驱动用于训练，它允许模型状态分布到可用GPU上。因此，每GPU内存消耗随着GPU数量增加而减少，允许DeepSpeed-HE支持更大的每GPU批量，导致超线性扩展。然而，在大规模上，虽然可用内存继续增加，但最大全局批量大小（在我们情况下为1024，序列长度为512）限制了每GPU批量大小，导致近线性或亚线性扩展。因此，对于给定的最大全局批量大小，DeepSpeed-HE在超线性和亚线性可扩展性的边界处实现最佳吞吐量和成本效率，确切点主要由可用内存和全局批量大小函数下的每GPU最大批量大小决定。

方法细节中的引用汇总。
- [6]：引用HuggingFace的Transformers，在描述易用体验和比较中，原文描述为“a pre-trained Huggingface [6] model”和“HuggingFace [6] powered by native PyTorch [14]”。展开为【[6], Huggingface’s transformers: State-of-the-art natural language processing+2019+arXiv preprint arXiv:1910.03771】。
- [7]：引用InstructGPT，在管道描述中，原文描述为“replicates the training pipeline from the InstructGPT [7] paper”和“highly consistent with InstructGPT paper [7]”。展开为【[7], Training language models to follow instructions with human feedback+2022+Advances in Neural Information Processing Systems】。
- [8]：引用RLHF，在管道中，原文描述为“Reinforcement Learning with Human Feedback (RLHF) [8]”。展开为【[8], Learning to summarize with human feedback+2020+Advances in Neural Information Processing Systems】。
- [9]：引用LoRA，在混合引擎中，原文描述为“LoRA [9]-based memory optimization strategies”和“Low Rank Adaptation (LoRA)”。展开为【[9], Lora: Low-rank adaptation of large language models+2021+arXiv preprint arXiv:2106.09685】。
- [11]：引用PPO，在步骤3中，原文描述为“using the Proximal Policy Optimization (PPO) [11] algorithm”。展开为【[11], Proximal policy optimization algorithms+2017+arXiv preprint arXiv:1707.06347】。
- [12]：引用SQuAD2.0，在混合训练中，原文描述为“prevent regression performance on public benchmarks like SQuAD2.0 [12]”。展开为【[12], Know what you don’t know: Unanswerable questions for squad+2018+arXiv preprint arXiv:1806.03822】。
- [13]：引用Colossal-AI，在比较中，原文描述为“Colossal-AI [13]”。展开为【[13], Colossal ai+2022+https://github.com/hpcaitech/ColossalAI】。
- [14]：引用PyTorch，在比较中，原文描述为“native PyTorch [14]”。展开为【[14], Pytorch: An imperative style, high-performance deep learning library+2019+Advances in neural information processing systems】。
- [15]：引用ZeRO，在可扩展性分析中，原文描述为“powered by ZeRO-based technology [15]”。展开为【[15], Zero: Memory optimizations toward training trillion parameter models+2020+SC20: International Conference for High Performance Computing, Networking, Storage and Analysis】。

实验环境

数据集名称、规模及用途。 使用多个数据集进行训练，支持数据抽象和混合，包括人类对查询的响应用于SFT、人类提供的排名数据集用于奖励模型微调，以及用于RLHF的反馈数据。支持如OPT系列模型的预训练数据，并提供数据分割/混合以在三个阶段中使用多个数据源。规模未指定精确数字，但支持从1.3B到175B参数模型的训练。

模型架构关键参数。 使用OPT系列模型，如OPT-1.3B、OPT-6.7B、OPT-13B、OPT-30B、OPT-66B、OPT-175B作为actor模型，OPT-350M作为奖励模型。序列长度为512，最大全局批量大小为1024。使用PPO算法，生成阶段初始prompt为256令牌，生成256令牌。

硬件配置。 GPU型号/数量/平台：NVIDIA A100-40G/80G（单节点8 GPU DGX、多节点如64 GPU）、NVIDIA A6000（单GPU，48GB内存）。网卡型号/数量未指定。CPU架构/型号未指定。连接关系：多节点多GPU集群，支持张量并行和ZeRO分片。

软件配置。 代码实现/语言：Python，使用DeepSpeed >=0.9.0、HuggingFace Transformers。依赖代码库：DeepSpeedExamples仓库、requirements.txt。OS未指定。

实验结果

训练时间和成本实验。 实验内容：在Azure上使用单节点8x A100评估不同模型的训练时间和成本。结果：OPT-13B需9小时、成本<300美元；OPT-30B需18小时、成本<600美元（表1）。分析结论：DeepSpeed-HE比现有系统快15倍，使RLHF训练快速且经济。

多节点训练时间和成本实验。 实验内容：在Azure上使用多节点64x A100-80GB评估训练时间和成本。结果：13B模型需1.25小时，175B模型需<1天（表2）。分析结论：展示优秀可扩展性，支持数百亿参数模型。

单GPU最大模型大小实验。 实验内容：评估单GPU支持的最大模型大小。结果：超过13亿参数（表3）。分析结论：使无多GPU访问的数据科学家能够训练大型模型，实现RLHF民主化。

吞吐量比较实验（单GPU）。 实验内容：与Colossal-AI和HuggingFace比较步骤3吞吐量，在单A100-40G GPU上。结果：DeepSpeed实现>10倍改进（图3）。分析结论：无图标表示OOM场景，DeepSpeed支持更大模型。

端到端训练吞吐量比较（多GPU）。 实验内容：在单DGX节点（8 A100-40G GPU）上比较不同模型大小的步骤3吞吐量。结果：DeepSpeed比Colossal-AI快6-19倍，比HuggingFace DDP快1.4-10.5倍（图4）。分析结论：解锁更大模型或更低成本训练。

生成阶段时间分解实验。 实验内容：分析1.3B模型在RLHF迭代中的时间/序列分解，在8 A100-40G GPU上。结果：生成阶段占大部分时间，DeepSpeed实现9x（vs HuggingFace）和15x（vs Colossal-AI）吞吐量改进（图5）。分析结论：混合引擎加速生成阶段，实现无与伦比效率。

有效吞吐量分析实验。 实验内容：评估不同模型大小的有效吞吐量（TFlops/GPU），在最大化效率的GPU数量下。结果：6.7B-66B模型最有效，175B比1.3B高1.2倍（图6）。分析结论：优化批量大小和TP，生成阶段占20%计算但可能主导时间。

可扩展性实验。 实验内容：评估13B和66B模型在增加DGX节点（8 A100-40/80G GPU）上的扩展。结果：整体良好扩展，小规模超线性、大规模近线性/亚线性（图7）。分析结论：ZeRO分区减少每GPU内存，允许更大批量；受最大全局批量1024限制。

模型规模比较实验。 实验内容：比较单GPU和单节点的最大模型大小。结果：DeepSpeed支持单GPU 6.5B、单节点50B，比Colossal-AI大7.5倍。分析结论：提升模型可扩展性。

结论

DeepSpeed-Chat使RLHF训练民主化，提供易用、高效且经济的端到端管道，支持从单GPU到多节点集群的数百亿参数模型训练。通过复制InstructGPT管道并统一DeepSpeed的训练和推理优化，它实现了超过现有系统的15倍速度、优秀可扩展性和单GPU上13B+模型支持，从而使资源有限的数据科学家能够创建强大模型，促进AI创新。未来工作包括持续改进DeepSpeed-Chat，根据反馈添加特征，如其路线图所示，并鼓励社区贡献和合作，以进一步发展DeepSpeed生态系统。