Microsoft DeepSpeed
要引用 DeepSpeed Chat，请引用我们的arxiv 报告：
@article{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}
}
ChatGPT 类模型席卷了人工智能领域，毫不夸张地说，它对数字世界的影响是革命性的。这些模型功能极其丰富，能够执行摘要、编码和翻译等任务，其结果堪比甚至超越人类专家。鉴于这些模型的强大功能，人工智能开源社区正在开展多项工作，以让 ChatGPT 类模型更容易被用户使用（例如 ChatLLaMa、Alpaca、Vicuna、Databricks-Dolly 等）。
尽管付出了这些令人难以置信的努力，但仍然缺乏一个端到端的 RLHF 流水线，无法训练出功能强大、易于 AI 社区使用的 ChatGPT 类模型。例如，使用现有系统训练一个 6.7B 的 ChatGPT 模型通常需要昂贵的多 GPU 配置，这超出了许多数据科学家的承受能力。即使拥有如此强大的计算资源，训练效率通常也不足这些机器自身能力的 5%（如下文所述）。最后，即使能够使用多 GPU 集群，现有解决方案也无法轻松、快速且经济地训练拥有数千亿参数的最先进的 ChatGPT 模型。
这些限制源于缺乏一个强大的系统设计，无法有效支持复杂的 InstructGPT RLHF 训练流程，这与现有深度学习系统所设计的标准预训练和微调流程截然不同。因此，本着普及 ChatGPT 类模型，并让 AI 社区真正能够使用 RLHF 训练的精神，我们今天发布了 DeepSpeed-Chat，它具有以下三个功能：
(i)易于使用的 ChatGPT 类模型训练和推理体验：只需一个脚本，即可使用预先训练好的 Huggingface 模型，使用 DeepSpeed-RLHF 系统运行 InstructGPT 训练的全部三个步骤，并生成您自己的 ChatGPT 类模型。此外，我们还提供推理 API，用于在模型训练完成后测试对话式交互。
（二）DeepSpeed-RLHF 流程：DeepSpeed-RLHF 流程主要复制了 InstructGPT 论文中的训练流程，并仔细确保其完整性以及与三个步骤一一对应，这三个步骤包括：a）监督微调 (SFT)，b）奖励模型微调，以及 c）基于人工反馈的强化学习 (RLHF)。此外，我们还提供数据抽象和混合功能，以便使用多个数据源进行训练。
(iii) DeepSpeed-RLHF 系统：一个强大而先进的 RLHF 系统，将 DeepSpeed 的训练和推理能力整合到一个统一的 RLHF 混合引擎 (DeepSpeed-HE) 中。该混合引擎能够在 RLHF 内部的推理和训练模式之间无缝切换，使其能够利用 DeepSpeed-Inference 的各种优化功能，例如张量并行和用于生成的高性能 Transformer 内核，同时还能受益于众多基于 ZeRO 和 LoRA 的内存优化策略，用于强化学习训练。DeepSpeed-HE 还能够感知完整的 RLHF 流水线，使其能够在 RLHF 不同阶段的内存管理和数据移动方面做出最佳决策。
DeepSpeed-RLHF 系统具有无与伦比的大规模效率，使复杂的 RLHF 训练变得快速、经济实惠，并且易于 AI 社区使用：
高效且经济实惠：在效率方面， DeepSpeed-HE 比现有系统快 15 倍以上，使 RLHF 训练既快速又经济实惠。例如，DeepSpeed-HE 在 Azure 云上仅需 9 小时即可训练 OPT-13B 模型，18 小时即可训练 OPT-30B 模型，前者的成本低于 300 美元，后者的成本低于 600 美元。
| GPU | OPT-6.7B | OPT-13B | OPT-30B | OPT-66B |
|---|---|---|---|---|
| 8个A100-40GB | 5.7小时 | 10.8小时 | 1.85天 | 不适用 |
| 8个A100-80GB | 4.1小时（132美元） | 9小时（290美元） | 18小时（580美元） | 2.1天（1620美元） |
表 1.单节点 8x A100：Azure 上的训练时间和相应的大约成本。
卓越的可扩展性：DeepSpeed-HE 支持数千亿参数的模型，并在多节点多 GPU 系统上实现卓越的可扩展性。即使是 13B 的模型，也能在 1.25 小时内完成训练，而使用 DeepSpeed-HE 训练 175B 的海量模型则可在一天内完成训练。
| GPU | OPT-13B | OPT-30B | OPT-66B | OPT-175B |
|---|---|---|---|---|
| 64个A100-80G | 1.25 小时（320 美元） | 4小时（1024美元） | 7.5小时（1920美元） | 20小时（5120美元） |
表 2.多节点 64x A100-80GB：Azure 上的训练时间和相应的大约成本。
非常重要的细节：以上两个表格中的数字均针对训练的第 3 步，并基于 DeepSpeed-RLHF 精选数据集和训练方案的实际测量训练吞吐量，该方案在总共 1.35 亿个 token 上训练了一个 epoch。我们总共拥有 67.5 亿个查询 token（131.9k 个序列长度为 256 的查询）和 67.5 亿个生成 token（131.9k 个序列长度为 256 的答案），每步的最大全局批处理大小为 0.5 亿个 token（1024 个查询-答案对）。我们强烈建议读者在与 DeepSpeed-RLHF 进行任何成本和端到端时间比较之前，务必注意这些规格。有关更多详细信息，请参阅我们的基准设置页面。
使 RLHF 训练民主化：仅使用一个 GPU，DeepSpeed-HE 就支持具有超过 130 亿个参数的训练模型，使无法使用多 GPU 系统的数据科学家不仅可以创建玩具 RLHF 模型，还可以创建可以在实际场景中使用的大型且功能强大的模型。
| V100 32G | A6000 48G | A100 40克 | A100 80克 | |
|---|---|---|---|---|
| 模型尺寸 | OPT-2.7B | OPT-6.7B | OPT-6.7B | OPT-13B |
表 3. DeepSpeed-HE 在单个 GPU 上支持的最大模型大小
接下来我们来深入探讨一下上面介绍的DeepSpeed-Chat的三种功能。
我们将从易于使用的体验开始，展示如何使用 DeepSpeed-RLHF 系统训练 OPT-13B 和 OPT-66B 模型。如果您时间紧迫，甚至可以在单个消费级 GPU 上仅用两个小时训练出 OPT-1.3B 模型。我们还将演示如何使用我们的 DeepSpeed-chat RLHF API 开发您自己的自定义流程。
a) 一个脚本完成 RLHF 训练的所有三个阶段并生成您的第一个 ChatGPT 模型！
我们在以下单个脚本中使用预训练的 OPT-13B 作为参与者模型，并使用 OPT-350M 作为奖励模型的示例，以生成最终的 13B ChatGPT 样式模型：
pip install deepspeed>=0.9.0
git clone https://github.com/deepspeedai/DeepSpeedExamples.git
cd DeepSpeedExamples/applications/DeepSpeed-Chat/
pip install -r requirements.txt
python train.py --actor-model facebook/opt-13b --reward-model facebook/opt-350m --deployment-type single_node
大约半天后，您的 13B 模型即可完成训练，并准备好检查点。下表列出了三个步骤的训练时间明细：
| 模型尺寸 | 步骤 1 | 第 2 步 | 步骤3 | 全部的 |
|---|---|---|---|---|
| 参与者：OPT-13B，奖励：OPT-350M | 2.5小时 | 0.25小时 | 10.8小时 | 13.6小时 |
表 4. 在配备 8 个 NVIDIA A100-40G GPU 的单个 DGX 节点上通过 DeepSpeed-Chat 训练 130 亿参数 ChatGPT 模型的 E2E 时间细分。
b) 使用 DeepSpeed-Chat 的推理 API 测试你的最终模型
现在是时候插入并测试由 DeepSpeed-Chat 训练的最终模型了！您不仅可以体验问答环节，还可以体验对话式的互动。
Human: Do you know Microsoft?
Assistant: Yes, I’m a big fan of Microsoft. What about them?
Human: Can you explain it to a 6-year old child? I wonder how I should describe it
Assistant: Sure, I can try. Microsoft is a company that makes computers, and they make a program called “Windows” which
is the operating system that runs on the computer. It’s like the “brain” of the computer. It’s where all the
programs and files are stored. It’s also where you’ll find the “start” button, which is where you’ll start to
use the computer.
我们了解用户通常喜欢尝试不同的模型大小和配置，以满足他们的训练时间、资源和质量要求。使用 DeepSpeed-Chat，用户可以轻松做到这一点。例如，如果您想在 GPU 集群上训练更大、更高质量的模型用于研究或业务，您只需使用相同的脚本，并根据所需的模型大小（例如 66B）和 GPU 数量（例如 64 个 GPU）进行设置：
python train.py --actor-model facebook/opt-66b --reward-model facebook/opt-350m --deployment-type multi_node
在 9 个小时内，您就可以准备好 660 亿个参数的 ChatGPT 模型，以便在您最喜欢的前端 GUI 中使用：
| 模型尺寸 | 步骤 1 | 第 2 步 | 步骤3 | 全部的 |
|---|---|---|---|---|
| 参与者：OPT-66B，奖励：OPT-350M | 82分钟 | 5分钟 | 7.5小时 | 9小时 |
表 5. 在 8 个 DGX 节点（每个节点配备 8 个 NVIDIA A100-80G GPU）上通过 DeepSpeed-Chat 训练 660 亿参数 ChatGPT 模型的 E2E 时间细分。
如果您只有 1-2 小时的咖啡或午休时间，您也可以尝试使用 DeepSpeed-Chat 训练一个小型/玩具模型。例如，我们准备了一个包含单个数据集的 13 亿模型的训练示例，以便在您的消费级 GPU 上测试我们的框架。最棒的是，当您午休回来时，您的模型检查点已经准备好了，可以直接使用！
python train.py --actor-model facebook/opt-1.3b --reward-model facebook/opt-350m --deployment-type single_gpu
| 模型尺寸 | 步骤 1 | 第 2 步 | 步骤3 | 全部的 |
|---|---|---|---|---|
| 参与者：OPT-1.3B，奖励：OPT-350M | 2900秒 | 670秒 | 1.2小时 | 2.2小时 |
表 6. 在具有 48GB 内存的单个商品 NVIDIA A6000 GPU 上通过 DeepSpeed-Chat 训练 13 亿参数 ChatGPT 模型的 E2E 时间细分。
DeepSpeed-Chat 允许用户使用我们如下所示的灵活 API 构建自己的 RLHF 训练流程，并可用于重构自己的 RLHF 训练策略。这提供了一个通用的界面和后端，可用于创建各种 RLHF 算法用于研究探索。
engine = DeepSpeedRLHFEngine(
actor_model_name_or_path=args.actor_model_name_or_path,
critic_model_name_or_path=args.critic_model_name_or_path,
tokenizer=tokenizer,
num_total_iters=num_total_iters,
args=args)
trainer = DeepSpeedPPOTrainer(engine=engine, args=args)
for prompt_batch in prompt_train_dataloader:
out = trainer.generate_experience(prompt_batch)
actor_loss, critic_loss = trainer.train_rlhf(out)
为了提供无缝的训练体验，我们遵循 InstructGPT，并在 DeepSpeed-Chat 中包含一个成熟的端到端训练流程，如图 1 所示。
图 1：DeepSpeed-Chat 的 RLHF 训练流程及其可选功能的图示。
我们的流程包括三个主要步骤：
- 步骤 1：监督微调（SFT），其中精心选择人类对各种查询的响应以微调预训练语言模型。
- 第 2 步：奖励模型微调，其中使用具有人类提供的对同一查询的多个答案的排名的数据集来训练单独的（通常比 SFT 更小）模型 (RW)。
- 步骤 3：RLHF 训练，其中使用近端策略优化 (PPO) 算法，利用来自 RW 模型的奖励反馈对 SFT 模型进行进一步微调。
我们在步骤 3 中提供了两个附加功能来帮助提高模型质量：
- 指数移动平均线 (EMA) 集合，其中可以选择基于 EMA 的检查点进行最终评估。
- 混合训练，将预训练目标（即下一个词预测）与 PPO 目标混合，以防止在 SQuAD2.0 等公共基准上出现回归性能。
EMA 和混合训练这两个训练功能，由于是可选的，经常被其他近期成果忽略。然而，根据 InstructGPT 的测试，EMA 检查点通常比传统的最终训练模型提供更好的响应质量，而混合训练可以帮助模型保留训练前的基准求解能力。因此，我们提供这些功能，以便用户充分获得 InstructGPT 中描述的训练体验，并追求更高的模型质量。
除了与 InstructGPT 论文高度一致之外，我们还提供了便捷的功能，支持研究人员和从业人员利用多种数据资源训练自己的 RLHF 模型：
- 数据抽象和混合功能： DeepSpeed-Chat 能够使用多个数据集训练模型，以提高模型质量。它具备以下功能：(1) 抽象数据集层，用于统一不同数据集的格式；(2) 数据拆分/混合功能，可将多个数据集适当地混合，然后拆分到 3 个训练阶段。
为了说明我们的训练流程的有效性，我们通过多轮对话来展示模型质量，如经验部分所示。
指令引导式 RLHF 流水线的步骤 1 和步骤 2 类似于大型模型的常规微调，它们由基于 ZeRO 的优化和 DeepSpeed 训练中灵活的并行策略组合提供支持，以实现规模和速度。另一方面，流水线的步骤 3 是性能影响方面最复杂的部分。每次迭代都需要高效处理两个阶段：a) 推理阶段，用于生成 token/经验，为训练提供输入；b) 训练阶段，用于更新参与者和奖励模型的权重，以及它们之间的交互和调度。这会带来两项主要成本：(1) 内存成本，因为在整个第 3 阶段需要提供 SFT 和 RW 模型的多个副本；(2) 主要的生成阶段，如果加速不当，将显著减慢整个第 3 阶段的速度。此外，我们在第 3 阶段添加的两个重要功能，包括指数移动平均线 (EMA) 收集和混合训练，将产生额外的内存和训练成本。
为了应对这些挑战，我们将 DeepSpeed 训练和推理的完整系统功能整合到一个统一的基础架构中，我们称之为混合引擎 (Hybrid Engine)。它利用原有的 DeepSpeed 引擎实现快速训练模式，同时轻松应用 DeepSpeed 推理引擎实现生成/评估模式，为第 3 阶段的 RLHF 训练提供显著更快的训练系统。如图 2 所示，DeepSpeed 训练和推理引擎之间的过渡是无缝的：通过为 Actor 模型启用典型的评估和训练模式，在运行推理和训练流水线时，DeepSpeed 会选择不同的优化方案来更快地运行模型并提高整体系统吞吐量。
图 2. DeepSpeed 混合引擎设计用于加速 RLHF 管道中最耗时的部分。
在 RLHF 训练的经验生成阶段的推理执行过程中，DeepSpeed Hybrid Engine 使用轻量级内存管理系统来处理 KV 缓存和中间结果，并结合高度优化的推理适配内核和张量并行实现，与现有解决方案相比，实现了吞吐量（每秒令牌数）的显著提升。
在训练执行过程中，混合引擎支持内存优化技术，例如 DeepSpeed 的 ZeRO 系列技术和低秩自适应 (LoRA)。我们设计和实现了这些系统优化，使它们彼此兼容，并可组合在一起，从而在统一的混合引擎下实现最高的训练效率。
混合引擎 (Hybrid Engine) 可以在训练和推理过程中无缝切换模型分区，以支持基于张量并行的推理和基于 ZeRO 的训练分片机制。它还可以重新配置内存系统，从而在每种模式下最大限度地提高内存可用性。这可以通过避免内存分配瓶颈并支持大批量大小来提升性能。混合引擎融合了 DeepSpeed 训练和推理的一系列系统技术，突破了现代 RLHF 训练的界限，并为 RLHF 工作负载提供了无与伦比的规模和系统效率。
如上所述，DeepSpeed-HE 是用于推理和训练的强大系统技术的融合，其架构旨在为各种硬件上的 DeepSpeed-RLHF 管道实现出色的规模和效率，从而使 RLHF 训练快速、经济实惠且易于 AI 社区使用。
在效率和可负担性方面，如表 1 所示，DeepSpeed-HE 在 Azure 云上仅需 9 小时即可训练 OPT-13B，并在 18 小时内训练 OPT-30B，费用分别低于 300 美元和 600 美元。在速度和可扩展性方面，如表 2 所示，即使是 13B 的模型也可以在 1.25 小时内完成训练，而使用 64 GPU 集群，则可在一天之内训练出 175B 的大型模型。在 RLHF 的可访问性和普及性方面，如表 3 所示，DeepSpeed-HE 支持在单个 GPU 上训练包含超过 130 亿个参数的模型。
与其他由原生 PyTorch 驱动的 RLHF 系统（如 Colossal-AI 或 HuggingFace）相比，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 倍。
因此，与现有的 RLHF 系统（如 Colossal-AI 或 HuggingFace DDP）相比，DeepSpeed-HE 的吞吐量提高了一个数量级，能够在相同的延迟预算下训练更大的参与者模型，或者以低 10 倍以上的成本训练类似大小的模型。
图 3.在单个 NVIDIA A100-40G 商用 GPU 上加速 RLHF 训练时，与另外两个系统框架进行的第 3 步吞吐量比较。图中无图标表示 OOM 场景。
图 4.在配备 8 块 NVIDIA A100-40G GPU 的单个 DGX 节点上，对训练流水线第 3 步（最耗时的部分）进行不同模型大小的端到端训练吞吐量比较。图
中无图标表示 OOM 场景。
这种效率的提升源于 DeepSpeed-HE 利用 DeepSpeed 推理优化技术加速 RLHF 处理中 RLHF 生成阶段的能力。图 5 展示了 13 亿参数模型在 RLHF 训练迭代中的时间分解：大部分时间都用于生成阶段。通过利用 DeepSpeed 的高性能推理内核，DeepSpeed-HE 在此阶段的吞吐量比 HuggingFace 提升高达 9 倍，比 Colossal-AI 提升高达 15 倍，从而实现无与伦比的端到端效率。
图 5.DeepSpeed Chat 混合引擎提供的卓越生成阶段加速：在配备 8 个 A100-40G GPU 的单个 DGX 节点上训练 OPT-1.3B 参与者模型 + OPT-350M 奖励模型的时间/序列细分。
（一）有效吞吐量分析。DeepSpeed -HE 在 RLHF 训练第三阶段的有效吞吐量取决于其在生成阶段和强化学习 (RL) 训练阶段实现的吞吐量。在我们的 RLHF 流水线中，生成阶段约占总计算量的 20%，而强化学习 (RL) 训练阶段则占剩余的 80%（详情请参阅基准设置页面）。然而，尽管前者所占比例较小，但它会占用大量的端到端 (e2e) 时间，因为它需要为 256 个生成的令牌中的每一个运行一次 Actor 模型，初始提示包含 256 个令牌，这使其受到内存带宽限制，难以实现高吞吐量。相比之下，强化学习 (RL) 训练阶段受计算限制，只需运行参考 Actor 模型进行几次前向和反向传递，每个样本包含来自提示和生成的全部 512 个令牌，即可实现良好的吞吐量。
图 6. 使用 DeepSpeed-HE 针对不同模型大小进行 RLHF 生成、训练和有效吞吐量，同时最大限度提高效率的 GPU 数量。
为了最大限度地提高有效吞吐量，DeepSpeed-HE 对两个阶段都进行了优化。首先，它使用尽可能大的批处理大小，以在两个阶段都获得更高的效率。其次，在生成阶段，当模型适合单个 GPU 内存时，它利用高性能 Transformer 内核来最大化 GPU 内存带宽利用率；当模型不适合单个 GPU 内存时，它利用张量并行 (TP)。在生成阶段使用 TP 而不是 ZeRO 来拟合模型，可以减少 GPU 之间的通信，并保持较高的 GPU 内存带宽利用率。
图 6 显示了 DeepSpeed-HE 在 1.3B 到 175B 规模的模型中，以 TFlops/GPU 计算时可实现的最佳有效吞吐量。图中还显示了每个生成和训练阶段所实现的吞吐量。对于 6.7B 到 66B 规模的模型，DeepSpeed-HE 的效率最高。当模型规模超过此范围至 175B 时，由于内存不足以支持更大的批次大小，吞吐量会有所下降，但仍然比 1.3B 的小型模型效率高出 1.2 倍。当我们将这些巨型模型扩展到更多 GPU，并为更大的批次大小提供更多内存时，它们的单 GPU 吞吐量可以进一步提升。
此外，我们想指出的是，如图 4 所示，我们的有效性能比现有系统高出 19 倍，这意味着它们的运行效率低于峰值的 5%。这不仅体现了优化 RLHF 工作负载的挑战性，也体现了我们系统在挑战面前依然保持的有效性。
图 7. 在配备 8 个 A100-40/80G GPU 的 DGX 节点数量不断增加的情况下，训练 13B（左）和 66B（右）Actor 模型 +350M 奖励模型的可扩展性
（II）可扩展性分析。不同模型大小的最佳有效吞吐量是在不同的 GPU 数量下实现的。部分原因是某些较大的模型需要更多内存才能运行。然而，这种现象很大程度上源于 DeepSpeed-HE 的可扩展性，我们将在下文讨论。
图 7 显示，DeepSpeed-RLHF 在多达 64 块 GPU 上实现了良好的整体扩展性。然而，如果我们更仔细地观察，就会发现 DeepSpeed-RLHF 训练在小规模下实现了超线性扩展，而在更大规模下则实现了近线性或亚线性扩展。这是由于内存可用性和最大全局批次大小之间的相互作用所致。
由于 DeepSpeed-HE 采用基于 ZeRO 的技术进行训练，它允许将模型状态划分到可用的 GPU 上。因此，每个 GPU 的内存消耗会随着 GPU 数量的增加而减少，从而使 DeepSpeed-HE 能够支持更大的单 GPU 批次，从而实现超线性扩展。然而，在大规模情况下，虽然可用内存持续增加，但最大全局批次大小（在我们的例子中为 1024，序列长度为 512）限制了单 GPU 的批次大小，导致近线性或亚线性扩展。因此，对于给定的最大全局批次大小，DeepSpeed-HE 在超线性和亚线性扩展性的边界上实现了最佳吞吐量和成本效率，而确切的临界点主要取决于单 GPU 可运行的最大批次大小，该大小与可用内存和全局批次大小有关。
我们非常高兴地分享 DeepSpeed-Chat 现已开源并可供 AI 社区使用。
-
首先，请访问 DeepSpeed-Chat 的 GitHub 页面：GitHub Landing Page
-
我们将根据您的反馈和支持，持续改进 DeepSpeed-Chat。我们的路线图展示了目前支持的功能以及未来计划的功能。
DeepSpeed-Chat 是 DeepSpeed 生态系统的一部分，该生态系统包含众多深度学习系统和建模技术。要了解更多信息，
- 请访问我们的网站以获取详细的博客文章、教程和有用的文档。
- 您还可以关注我们的英文 Twitter、日文 Twitter和中文知乎，以获取有关 DeepSpeed 的最新消息。
如果这篇文章帮助到了你，你可以请作者喝一杯咖啡
浙公网安备 33010602011771号