DeepSpeed-Chat
DeepSpeed-Chat: 轻松、快速且经济高效的RLHF训练，适用于所有规模的ChatGPT类模型:
摘要
ChatGPT类模型已经在人工智能的各个应用领域引发了革命，从摘要生成、编码到翻译，其表现甚至超越人类。然而，当前的AI领域缺乏一个易于使用、高效且经济实惠的端到端RLHF（基于人类反馈的强化学习）训练 pipeline，尤其是在训练数十亿参数规模的模型时。本文介绍了DeepSpeed-Chat，这是一个新颖的系统，旨在使RLHF训练民主化，使其对AI社区更加可及。DeepSpeed-Chat提供了三个关键功能：易于使用的ChatGPT类模型训练和推理体验、DeepSpeed-RLHF pipeline（复制了InstructGPT的训练 pipeline），以及一个强大的DeepSpeed-RLHF系统，该系统将各种优化技术统一结合，用于训练和推理。该系统提供了无与伦比的效率和可扩展性，能够在创纪录的时间内以低成本训练数百亿参数的模型。通过这一进展，DeepSpeed-Chat为更广泛的RLHF训练铺平了道路，即使对于资源有限的数据科学家也是如此，从而促进了AI领域的创新和进一步发展。
1 概述
ChatGPT [1] 类模型已经席卷了AI世界，毫不夸张地说，它对数字世界的影响是革命性的。这些模型非常多功能，能够执行摘要生成、编码和翻译等任务，其结果与人类专家相当甚至超越。鉴于这些模型的强大功能，AI开源社区正在努力使ChatGPT风格的模型更加易于使用（例如ChatLLaMa [2]、Alpaca [3]、Vicuna [4]、Databricks-Dolly [5]等）。
尽管这些努力令人惊叹，但仍然缺乏一个端到端的RLHF pipeline，能够训练强大的ChatGPT类模型，并且易于AI社区使用。例如，使用现有系统训练一个适中的6.7B ChatGPT模型通常需要昂贵的多GPU设置，这超出了许多数据科学家的能力范围。即使拥有这些计算资源，训练效率通常也低于这些机器能力的5%（如图6所示）。最后，现有解决方案根本无法支持轻松、快速且经济实惠地训练数百亿参数的最先进ChatGPT模型，即使拥有多GPU集群。
这些限制源于缺乏一个强大的系统设计，能够有效支持复杂的InstructGPT RLHF训练 pipeline，这与现有深度学习系统设计的标准预训练和微调 pipeline有很大不同。因此，本着使ChatGPT类模型民主化并使RLHF训练真正对AI社区可及的精神，今天我们发布了DeepSpeed-Chat，它具有以下三个功能：
- 易于使用的ChatGPT类模型训练和推理体验：一个简单的脚本，能够使用DeepSpeed-RLHF系统将预训练的Huggingface [6]模型通过InstructGPT [7]训练的三个步骤运行，并生成您自己的ChatGPT类模型。此外，我们提供了一个推理API，用于在模型训练后测试对话式交互。
- DeepSpeed-RLHF pipeline：DeepSpeed-RLHF pipeline主要复制了InstructGPT [7]论文中的训练 pipeline，并确保与三个步骤的完整性和一一对应，包括：a) 监督微调（SFT），b) 奖励模型微调，c) 基于人类反馈的强化学习（RLHF）[8]。此外，我们提供了数据抽象和混合功能，以支持使用多个数据源进行训练。
- DeepSpeed-RLHF系统：一个强大而复杂的RLHF系统，将DeepSpeed的训练和推理能力结合到一个统一的混合引擎（DeepSpeed-HE）中。混合引擎能够在RLHF中无缝切换推理和训练模式，使其能够利用DeepSpeed-Inference的各种优化，如张量并行和高性能Transformer内核生成，同时也能受益于基于ZeRO和LoRA [9]的内存优化策略。DeepSpeed-HE还了解整个RLHF pipeline，使其能够在RLHF的不同阶段做出最佳的内存管理和数据移动决策。
DeepSpeed-RLHF系统在规模上具有无与伦比的效率，使复杂的RLHF训练变得快速、经济实惠且易于AI社区使用：
效率和成本效益：在效率方面，DeepSpeed-HE比现有系统快15倍以上，使RLHF训练既快速又经济实惠。例如，DeepSpeed-HE可以在Azure云上在9小时内训练OPT-13B [10]，在18小时内训练OPT-30B，成本分别低于300美元和600美元，如表1所示。
出色的可扩展性：DeepSpeed-HE支持数百亿参数的模型，并可以在多节点多GPU系统上实现出色的可扩展性。因此，即使是13B模型也可以在1.25小时内训练完成，而巨大的175B模型可以在DeepSpeed-HE的帮助下在一天内训练完成，如表2所示。
表1：单节点8x A100：在Azure上的训练时间和相应的近似成本
| GPUs | OPT-6.7B | OPT-13B | OPT-30B | OPT-66B |
|---|---|---|---|---|
| 8x A100-40GB | 5.7小时 | 10.8小时 | 1.85天 | 不适用 |
| 8x A100-80GB | 4.1小时 ($132) | 9小时 ($290) | 18小时 ($580) | 2.1天 ($1620) |
表2：多节点64x A100-80GB：在Azure上的训练时间和相应的近似成本
| GPUs | OPT-13B | OPT-30B | OPT-66B | OPT-175B |
|---|---|---|---|---|
| 64x A100-80G | 1.25小时 ($320) | 4小时 ($1024) | 7.5小时 ($1920) | 20小时 ($5120) |
表3：DeepSpeed-HE在单个GPU上支持的最大模型大小
| 模型大小 | OPT-2.7B | OPT-6.7B | OPT-6.7B | OPT-13B |
|---|
表4：在单个DGX节点上使用8个NVIDIA A100-40G GPU通过DeepSpeed-Chat训练130亿参数ChatGPT模型的端到端时间分解
| 模型大小 | 步骤1 | 步骤2 | 步骤3 | 总计 |
|---|---|---|---|---|
| 演员：OPT-13B，奖励：OPT-350M | 2.5小时 | 0.25小时 | 10.8小时 | 13.6小时 |
表5：在8个DGX节点上使用8个NVIDIA A100-80G GPU/节点通过DeepSpeed-Chat训练660亿参数ChatGPT模型的端到端时间分解
| 模型大小 | 步骤1 | 步骤2 | 步骤3 | 总计 |
|---|---|---|---|---|
| 演员：OPT-66B，奖励：OPT-350M | 82分钟 | 5分钟 | 7.5小时 | 9小时 |
表6：在单个消费级NVIDIA A6000 GPU（48GB内存）上通过DeepSpeed-Chat训练13亿参数ChatGPT模型的端到端时间分解
| 模型大小 | 步骤1 | 步骤2 | 步骤3 | 总计 |
|---|---|---|---|---|
| 演员：OPT-1.3B，奖励：OPT-350M | 2900秒 | 670秒 | 1.2小时 | 2.2小时 |
RLHF训练的民主化：仅使用单个GPU，DeepSpeed-HE支持训练超过130亿参数的模型，如表3所示，使没有多GPU系统的数据科学家能够创建不仅仅是玩具RLHF模型，而是可以用于实际场景的大型强大模型。
接下来，我们将深入介绍上述DeepSpeed-Chat的三个功能。
2 易于使用的ChatGPT训练和推理体验
我们首先展示如何使用DeepSpeed-RLHF系统训练OPT-13B和OPT-66B模型。如果您时间紧迫，您甚至可以在单个消费级GPU上在短短两小时内训练一个OPT-1.3B模型。我们还展示了如何使用我们的DeepSpeed-chat RLHF API开发您自己的自定义 pipeline。
2.1 使用DeepSpeed-Chat的RLHF示例训练您的第一个ChatGPT风格模型
一个简单的脚本即可完成RLHF训练的所有三个阶段，并生成您的第一个ChatGPT模型！我们使用预训练的OPT-13B作为演员模型，OPT-350M作为奖励模型，在以下脚本中生成最终的13B ChatGPT风格模型：
大约半天时间，您的13B模型将完全训练完成，检查点准备就绪。表4展示了每个训练阶段的时间分解。
使用DeepSpeed-Chat的推理API测试您的最终模型现在可以插入并测试由DeepSpeed-Chat训练的最终模型！您不仅可以进行问答会话，还可以进行对话式交互。
人类：你知道微软吗？ 助手：是的，我是微软的忠实粉丝。他们怎么样？
人类：你能向一个6岁的孩子解释吗？我想知道该如何描述它
助手：当然，我可以试试。微软是一家制造计算机的公司，他们制作了一个名为“Windows”的程序，这是运行在计算机上的操作系统。它就像计算机的“大脑”，所有程序和文件都存储在那里。你还可以在那里找到“开始”按钮，这是你开始使用计算机的地方。
2.2 想尝试不同的模型大小和配置？没问题！
我们理解用户通常喜欢尝试不同的模型大小和配置，以满足他们的训练时间、资源和质量要求。使用DeepSpeed-Chat，用户可以轻松做到这一点。例如，如果您想在GPU集群上训练一个更大、更高质量的模型用于研究或业务，您可以简单地使用相同的脚本，指定所需的模型大小（例如66B）和GPU数量（例如64个GPU）：
在9小时内，您就可以准备好660亿参数的ChatGPT模型，可以在您喜欢的前端GUI中提供服务，如表5所示。
如果您只有大约1-2小时的咖啡或午餐时间，您也可以尝试使用DeepSpeed-Chat训练一个小型/玩具模型。例如，我们准备了一个1.3B模型的训练示例，使用单个数据集在您的消费级GPU上测试我们的框架。最棒的是，当您从午餐休息回来时，您的模型检查点已经准备好可以使用了（见表6）！
2.3 使用DeepSpeed-Chat的RLHF API自定义您自己的RLHF训练 pipeline
DeepSpeed-Chat允许用户使用我们灵活的API构建他们自己的RLHF训练 pipeline，如下所示，用户可以使用这些API重建他们自己的RLHF训练策略。这为创建各种RLHF算法提供了一个通用接口和后端，用于研究探索。
3 完整的RLHF训练 pipeline
为了提供无缝的训练体验，我们遵循InstructGPT，并在DeepSpeed-Chat中包含了一个完整的端到端训练 pipeline，如图1所示。
我们的 pipeline包括三个主要步骤：
步骤1：监督微调（SFT），其中人类对各种查询的响应被精心选择，用于微调预训练的语言模型。
步骤2：奖励模型微调，其中使用一个单独（通常比SFT模型小）的模型（RW）训练一个数据集，该数据集包含人类提供的对同一查询的多个答案的排名。
步骤3：RLHF训练，其中SFT模型使用RW模型的奖励反馈通过近端策略优化（PPO）[11]算法进一步微调。
我们在步骤3中提供了两个额外的功能，以帮助提高模型质量：
指数移动平均（EMA）收集，其中可以选择一个基于EMA的检查点进行最终评估。
混合训练，它将预训练目标（即下一个单词预测）与PPO目标混合，以防止在公共基准测试（如SQuAD2.0 [12]）上的回归性能。
这两个训练功能，EMA和混合训练，通常被其他最近的努力忽略，因为它们可以是可选的。然而，根据InstructGPT，EMA检查点通常比传统的最终训练模型提供更好的响应质量，而混合训练可以帮助模型保留预训练基准测试的解决能力。因此，我们为用户提供了这些功能，以完全按照InstructGPT的描述获得训练体验，并争取更高的模型质量。
除了与InstructGPT论文[7]高度一致外，我们还提供了方便的功能，支持研究人员和从业者使用多个数据资源训练他们自己的RLHF模型：
- 数据抽象和混合功能：DeepSpeed-Chat能够使用多个数据集训练模型，以提高模型质量。它配备了（1）一个抽象的数据集层，以统一不同数据集的格式；（2）数据拆分/混合功能，以便多个数据集在三个训练阶段中正确混合和拆分。
为了说明我们训练 pipeline的有效性，我们在体验部分展示了多轮对话的模型质量。
图1：DeepSpeed-Chat的RLHF训练 pipeline示意图，包含可选功能。
4 DeepSpeed混合引擎 - 统一的基础设施为RLHF训练提供动力和优化
InstructGPT RLHF pipeline的步骤1和步骤2类似于大型模型的常规微调，它们由基于ZeRO的优化和DeepSpeed训练中的灵活并行策略组合提供动力，以实现规模和速度。另一方面， pipeline的步骤3在性能影响方面是最复杂的部分。每次迭代都需要高效处理两个阶段：a) 推理阶段用于生成令牌/经验，为训练生成输入；b) 训练阶段用于更新演员和奖励模型的权重，以及它们之间的交互和调度。它引入了两个主要成本：（1）内存成本，因为在阶段3中需要提供SFT和RW模型的多个副本；（2）主要的生成阶段，如果没有适当加速，将显著减慢整个阶段3。此外，我们在阶段3中添加的两个重要功能，包括指数移动平均（EMA）收集和混合训练，将产生额外的内存和训练成本。
为了应对这些挑战，我们将DeepSpeed训练和推理的完整系统能力组合到一个统一的基础设施中，我们称之为混合引擎。它利用原始的DeepSpeed引擎进行快速训练模式，同时轻松应用DeepSpeed推理引擎进行生成/评估模式，为RLHF训练在阶段3提供了一个显著更快的训练系统。如图2所示，DeepSpeed训练和推理引擎之间的过渡是无缝的：通过为演员模型启用典型的评估和训练模式，在运行推理和训练 pipeline时，DeepSpeed选择其不同的优化以更快地运行模型并提高整体系统吞吐量。
在RLHF训练的经验生成阶段的推理执行期间，DeepSpeed混合引擎使用轻量级内存管理系统处理KV缓存和中间结果，结合高度优化的推理适配内核和张量并行实现，与现有解决方案相比，显著提高了吞吐量（每秒令牌数）。
在训练执行期间，混合引擎启用了内存优化技术，如DeepSpeed的ZeRO系列技术和低秩适应（LoRA）。我们设计和实现了这些系统优化，使它们彼此兼容，并可以组合在一起，在统一的混合引擎下提供最高的训练效率。
混合引擎可以无缝地在训练和推理之间更改模型分区，以支持基于张量并行的推理和基于ZeRO的分片机制进行训练。它还可以重新配置内存系统，以在每个模式下最大化内存可用性。这通过避免内存分配瓶颈和支持大批量大小来提高性能。混合引擎集成了DeepSpeed训练和推理的一系列系统技术，推动了现代RLHF训练的边界，并为RLHF工作负载提供了无与伦比的规模和系统效率。
5 DeepSpeed RLHF：通过混合引擎实现无与伦比的规模和效率
5.1 功能回顾
如前所述，DeepSpeed-HE是推理和训练的强大系统技术的结合，旨在为DeepSpeed-RLHF pipeline在各种硬件上实现出色的规模和效率，使RLHF训练快速、经济实惠且易于AI社区使用。
在效率和成本效益方面，如表1所示，DeepSpeed-HE可以在Azure云上在9小时内训练OPT-13B，在18小时内训练OPT-30B，成本分别低于300美元和600美元。在速度和可扩展性方面，如表2所示，即使是13B模型也可以在1.25小时内训练完成，而巨大的175B模型可以在64 GPU集群的帮助下在一天内训练完成。在RLHF的民主化和可及性方面，DeepSpeed-HE支持在单个GPU上训练超过130亿参数的模型，如表3所示。
5.2 与现有RLHF系统的吞吐量和模型大小可扩展性比较
与其他RLHF系统（如Colossal-AI [13]或由原生PyTorch [14]驱动的HuggingFace [6]）相比，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解锁了在相同延迟预算下训练显著更大的演员模型或以超过10倍更低的成本训练相似大小模型的能力，与现有的RLHF系统（如Colossal-AI或HuggingFace DDP）相比。
这种效率的提升源于DeepSpeed-HE能够利用DeepSpeed推理优化加速RLHF生成阶段。图5显示了在RLHF训练迭代中1.3B参数模型的时间分解：大部分时间用于生成阶段。通过利用DeepSpeed的高性能推理内核，DeepSpeed-HE在该阶段可以实现比HuggingFace高9倍的吞吐量，比Colossal-AI高15倍，从而实现无与伦比的端到端效率。
5.3 有效吞吐量和可扩展性分析
(I) 有效吞吐量分析。DeepSpeed-HE在RLHF训练阶段3的有效吞吐量取决于它在生成和RL训练阶段实现的吞吐量。在我们的RLHF pipeline中，生成阶段约占总计算的20%，而RL训练阶段占剩余的80%（详见基准设置https://github.com/microsoft/DeepSpeedExamples/tree/master/applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/BenchmarkSetting.md）。然而，尽管生成阶段占比较小，但它可能占据端到端时间的很大一部分，因为它需要为每个生成的256个令牌运行一次演员模型，初始提示为256个令牌，这使得它受限于内存带宽，难以实现高吞吐量。相比之下，RL训练阶段是计算密集型的，运行参考演员模型时只需几次前向和后向传递，每个样本包含完整的512个令牌（提示和生成），因此可以实现良好的吞吐量。
为了最大化有效吞吐量，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显示，DeepSeed-RLHF在最多64个GPU上实现了良好的整体扩展。然而，如果我们更仔细地观察，可以发现DeepSpeed-RLHF训练在小规模时实现了超线性扩展，随后在大规模时实现了接近线性或次线性扩展。这是由于内存可用性和最大全局批量大小之间的相互作用。
由于DeepSpeed-HE由基于ZeRO的技术[15]提供支持进行训练，它允许模型状态在可用GPU之间进行分区。因此，随着GPU数量的增加，每个GPU的内存消耗减少，使DeepSpeed-HE能够支持每个GPU的更大批量，从而实现超线性扩展。然而，在大规模时，虽然可用内存继续增加，但最大全局批量大小（在我们的情况下为1024，序列长度为512）限制了每个GPU的批量大小，导致接近线性或次线性扩展。因此，对于给定的最大全局批量大小，DeepSpeed-HE在超线性和次线性可扩展性的边界处实现了最佳吞吐量和成本效率，而确切点主要由每个GPU可以运行的最大批量大小决定，作为可用内存和全局批量大小的函数。