RLHF 训练框架 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 系统将预训练的 Hugging Face [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) 推理阶段用于生成 Token/经验，为训练生成输入； b) 训练阶段用于更新演员和奖励模型的权重，以及它们之间的交互和调度。它引入了两个主要成本：（1）内存成本，因为在阶段 3 中需要提供 SFT 和 RW 模型的多个副本；（2）主要的生成阶段，如果没有适当加速，将显著减慢整个阶段 3。此外，我们在阶段 3 中添加的两个重要功能，包括指数移动平均（EMA）收集和混合训练，将产生额外的内存和训练成本。
为了应对这些挑战，我们将 DeepSpeed 训练和推理的完整系统能力组合到一个统一的基础设施中，我们称之为混合引擎。它利用原始的 DeepSpeed 引擎进行快速训练模式，同时轻松应用 DeepSpeed 推理引擎进行生成/评估模式，为 RLHF 训练在阶段 3 提供了一个显著更快的训练系统。如图 2 所示， DeepSpeed 训练和推理引擎之间的过渡是无缝的：通过为演员模型启用典型的评估和训练模式，在运行推理和训练 pipeline 时， DeepSpeed 选择其不同的优化以更快地运行模型并提高整体系统吞吐量。
在 RLHF 训练的经验生成阶段的推理执行期间， DeepSpeed 混合引擎使用轻量级内存管理系统处理 KV 缓存和中间结果，结合高度优化的推理适配内核和张量并行实现，与现有解决方案相比，显著提高了吞吐量（每秒 Token 数）。
在训练执行期间，混合引擎启用了内存优化技术，如 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]驱动的 Hugging Face [6]）相比， DeepSpeed-RLHF 在系统性能和模型可扩展性方面表现出色：
在吞吐量方面， DeepSpeed 在单个 GPU 上的 RLHF 训练中实现了超过 10 倍的改进（图 3）。在多 GPU 设置中，它比 Colossal-AI 实现了 6-19 倍的加速，比 Hugging Face 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 或 Hugging Face DDP）相比。
这种效率的提升源于 DeepSpeed-HE 能够利用 DeepSpeed 推理优化加速 RLHF 生成阶段。图 5 显示了在 RLHF 训练迭代中 1.3B 参数模型的时间分解：大部分时间用于生成阶段。通过利用 DeepSpeed 的高性能推理内核， DeepSpeed-HE 在该阶段可以实现比 Hugging Face 高 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 个 Token 运行一次演员模型，初始提示为 256 个 Token，这使得它受限于内存带宽，难以实现高吞吐量。相比之下， RL 训练阶段是计算密集型的，运行参考演员模型时只需几次前向和后向传递，每个样本包含完整的 512 个 Token（提示和生成），因此可以实现良好的吞吐量。
为了最大化有效吞吐量， 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 可以运行的最大批量大小决定，作为可用内存和全局批量大小的函数。