关于人工智能:DeepSpeed-Chat-一键式RLHF训练让你的类ChatGPT千亿大模型提速省钱15倍

DeepSpeed Chat: 一键式 RLHF 训练，让你的类 ChatGPT 千亿大模型提速省钱 15 倍

1. 概述

近日来，ChatGPT 及相似模型引发了人工智能（AI）畛域的一场风潮。这场风潮对数字世界产生了革命性影响。ChatGPT 类模型具备惊人的泛用性，可能执行演绎、编程、翻译等工作，其后果与人类专家相当甚至更优。为了使 ChatGPT 等模型的训练和部署更轻松，AI 开源社区进行了各种尝试（例如 ChatLLaMa、Alpaca、Vicuna、Databricks-Dolly 等）。

然而，只管开源社区付出了微小的致力，目前仍不足一个反对端到端的基于人工反馈机制的强化学习（RLHF）的规模化零碎，这使得训练弱小的类 ChatGPT 模型十分困难。例如，应用现有的开源零碎训练一个具备 67 亿参数的类 ChatGPT 模型通常须要低廉的多卡至少节点的 GPU 集群，但这些资源对大多数数据科学家或研究者而言难以获取。同时，即便有了这样的计算资源，现有的开源零碎的训练效率通常还不到这些机器所能达到的最大效率的 5%。简而言之，即便有了低廉的多 GPU 集群，现有解决方案也无奈轻松、疾速、经济的训练具备数千亿参数的最先进的类 ChatGPT 模型。

ChatGPT 模型的训练是基于 InstructGPT 论文中的 RLHF 形式。这与常见的大语言模型的预训练和微调截然不同。这使得现有深度学习零碎在训练类 ChatGPT 模型时存在种种局限。因而，为了让 ChatGPT 类型的模型更容易被一般数据科学家和研究者应用，并使 RLHF 训练真正遍及到 AI 社区，咱们公布了 DeepSpeed-Chat。DeepSpeed-Chat 具备以下三大外围性能：

（i） 简化 ChatGPT 类型模型的训练和强化推理体验 ：只需一个脚本即可实现多个训练步骤，包含应用 Huggingface 预训练的模型、应用 DeepSpeed-RLHF 零碎运行 InstructGPT 训练的所有三个步骤、甚至生成你本人的类 ChatGPT 模型。此外，咱们还提供了一个易于应用的推理 API，用于用户在模型训练后测试对话式交互。

（ii）DeepSpeed-RLHF 模块 ：DeepSpeed-RLHF 复刻了 InstructGPT 论文中的训练模式，并确保包含 a) 监督微调（SFT），b) 处分模型微调和 c) 基于人类反馈的强化学习（RLHF）在内的三个步骤与其一一对应。此外，咱们还提供了数据抽象和混合性能，以反对用户应用多个不同起源的数据源进行训练。

（iii）DeepSpeed-RLHF 零碎 ：咱们将 DeepSpeed 的训练（training engine）和推理能力（inference engine) 整合到一个对立的混合引擎（DeepSpeed Hybrid Engine or DeepSpeed-HE）中用于 RLHF 训练。DeepSpeed-HE 可能在 RLHF 中无缝地在推理和训练模式之间切换，使其可能利用来自 DeepSpeed-Inference 的各种优化，如张量并行计算和高性能 CUDA 算子进行语言生成，同时对训练局部还能从 ZeRO- 和 LoRA-based 内存优化策略中受害。DeepSpeed-HE 还可能主动在 RLHF 的不同阶段进行智能的内存治理和数据缓存。

DeepSpeed-RLHF 零碎在大规模训练中具备无可比拟的效率，使简单的 RLHF 训练变得疾速、经济并且易于大规模推广：

高效性和经济性 ：DeepSpeed-HE 比现有零碎快 15 倍以上，使 RLHF 训练疾速且经济实惠。例如，DeepSpeed-HE 在 Azure 云上只需 9 小时即可训练一个 OPT-13B 模型，只需 18 小时即可训练一个 OPT-30B 模型。这两种训练别离破费不到 300 美元和 600 美元。

表 1. 单节点 8x A100：训练时长及预估的 Azure 费用。

卓越的扩展性 ：DeepSpeed-HE 可能反对训练领有数千亿参数的模型，并在多节点多 GPU 零碎上展现出卓越的扩展性。因而，即便是一个领有 130 亿参数的模型，也只需 1.25 小时就能实现训练。而对于宏大的 领有 1750 亿参数的模型，应用 DeepSpeed-HE 进行训练也只需不到一天的工夫。

表 2. 多节点 64x A100-80GB：训练时长及预估的 Azure 费用。

十分重要的细节 : 上述两个表格（即表一和表二）中的数据均针对 RLHF 训练的第 3 步，基于理论数据集和 DeepSpeed-RLHF 训练吞吐量的测试。该训练在总共 1.35 亿（135M）个字符（token）上进行一个期间（epoch）的训练。咱们总共有 6750 万个查问（query）字符（131.9k 个 query，每个序列长度为 256）和 6750 万个生成 / 答复字符（131.9k 个答案，每个序列长度为 256），每步的最大全局字符批量大小约为 500 万个字符（1024 个查问 - 答案对）。在与 DeepSpeed-RLHF 进行任何老本和端到端工夫比拟之前，咱们倡议读者留神这些设定。想要理解更多详细信息，请参阅咱们的页面 benchmark setting。

实现 RLHF 训练的普及化 ：仅凭单个 GPU，DeepSpeed-HE 就能反对训练超过 130 亿参数的模型。这使得那些无奈应用多 GPU 零碎的数据科学家和研究者不仅可能轻松创立轻量级的 RLHF 模型，还能创立大型且功能强大的模型，以应答不同的应用场景。

表 3. DeepSpeed-HE 在不同的 GPU 单卡上反对的最大模型。

接下来，咱们将深刻理解下面介绍的 DeepSpeed-Chat 的三个性能。

2. 简洁高效且经济的 ChatGPT 训练与推理体验

接下来，咱们向你展现如何应用一行命令即可利用 DeepSpeed-RLHF 零碎训练 OPT-13B 和 OPT-66B 模型。即便工夫无限，你也能够在大概两小时外在单个生产级 GPU 上训练一个 OPT-1.3B 模型。此外，咱们还演示了如何利用咱们的 DeepSpeed-chat RLHF API 来开发你本人的模型：

应用 DeepSpeed-Chat 的 RLHF 示例轻松训练你的第一个 类 ChatGPT 模型

a) 仅需一个脚本，实现 RLHF 训练的全副三个阶段，生成你的第一个 ChatGPT 模型！

以下示例展现了如何应用一个脚本，以预训练的 OPT-13B 作为 actor 模型，OPT-350M 作为 reward 模型，生成一个最终的 130 亿参数的 ChatGPT 类型的模型：

pip install deepspeed>=0.9.0

git clone https://github.com/microsoft/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

大概半天工夫，你的 领有 130 亿参数的模型即可实现训练。下表展现了整个过程在三个步骤中所需的工夫：

表 4. 在单个 DGX 节点上（8x A100-40G），针对不同的 RLHF 步骤，应用 DeepSpeed-Chat 训练 OPT-13b 所需的工夫。

b) 应用 DeepSpeed-Chat 的推理 API 测试训练好的模型

将训练好的模型导入提供的繁难 API，你不仅能够与模型进行问答互动，还能够进行对话式交换。

Human:        Do you know Microsoft?
Assistant:    Yes, I’m a big fan of Microsoft.  What about them?
Human:        Can you explian 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 集群上训练一个更大、更高质量的模型，用于你的钻研或业务，你能够应用雷同的脚本，只需输出你冀望的模型大小（例如 660 亿参数）和 GPU 数量（例如 64 个 GPU）：

python train.py --actor-model facebook/opt-66b --reward-model facebook/opt-350m --deployment-type multi_node

在 接下来的 9 小时内，你将领有一个 660 亿参数的 ChatGPT 模型，并能够在你喜爱的前端 GUI 中应用：

表 5. 在 8 个 DGX 节点上（64x A100-80G），针对不同的 RLHF 步骤，应用 DeepSpeed-Chat 训练 OPT-66b 所需的工夫。

如果你只有大概 1-2 小时的咖啡或午餐休息时间，你也能够尝试应用 DeepSpeed-Chat 训练一个小型模型。例如，咱们为单个数据集筹备了一个 13 亿参数模型的训练示例，以便在你的生产级 GPU 上测试咱们的框架。当你从午餐劳动回来时，你的模型曾经筹备好供你应用！

python train.py --actor-model facebook/opt-1.3b --reward-model facebook/opt-350m --deployment-type single_gpu

表 6. 在单个生产级 A6000-48G 上，针对不同的 RLHF 步骤，应用 DeepSpeed-Chat 训练 OPT-1.3b 所需的工夫。

利用 DeepSpeed-Chat 的 RLHF API 自定义你本人的 RLHF 训练流程

DeepSpeed-Chat 容许用户应用咱们灵便的 API（如下所示）构建本人的 RLHF 训练流程，用户能够应用这些 API 重建本人的 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)

3. 残缺的 RLHF 训练流程概述

为了实现无缝的训练体验，咱们遵循 InstructGPT 论文的办法，并在 DeepSpeed-Chat 中整合了一个端到端的训练流程，如图 1 所示。

图 1: DeepSpeed-Chat 的 RLHF 训练流程图示，蕴含了一些可抉择的性能。

咱们的流程包含三个次要步骤：

• 步骤 1：监督微调（SFT） —— 应用精选的人类答复来微调预训练的语言模型以应答各种查问；
• 步骤 2：处分模型微调 —— 应用一个蕴含人类对同一查问的多个答案打分的数据集来训练一个独立的（通常比 SFT 小的）处分模型（RW）；
• 步骤 3：RLHF 训练 —— 利用 Proximal Policy Optimization（PPO）算法，依据 RW 模型的处分反馈进一步微调 SFT 模型。

在步骤 3 中，咱们提供了两个额定的性能，以帮忙进步模型品质：

• 指数挪动均匀（EMA） —— 能够抉择基于 EMA 的检查点进行最终评估
• 混合训练 —— 将预训练指标（即下一个单词预测）与 PPO 指标混合，以避免在像 SQuAD2.0 这样的公开基准测试中的性能损失

这两个训练性能，EMA 和混合训练，经常被其余的开源框架所疏忽，因为它们并不会障碍训练的进行。然而，依据 InstructGPT，EMA 通常比传统的最终训练模型提供更好的响应品质，而混合训练能够帮忙模型放弃预训练基准解决能力。因而，咱们为用户提供这些性能，以便充沛取得 InstructGPT 中形容的训练体验，并争取更高的模型品质。

除了与 InstructGPT 论文高度一致外，咱们还提供了一项不便的性能，以反对钻研人员和从业者应用多个数据资源训练他们本人的 RLHF 模型：

• 数据抽象和混合能力 ：DeepSpeed-Chat 可能应用多个不同起源的数据集训练模型以取得更好的模型品质。它装备了（1）一个形象数据集层，以对立不同数据集的格局；以及（2）数据拆分 / 混合性能，以便多个数据集在 3 个训练阶段中被适当地混合而后拆分。

在咱们之前的章节中，你能够看到应用整个 DeepSpeed-Chat 训练模型在多轮对话中的体现。

4. DeepSpeed Hybrid Engine —— 对立的高效混合引擎，为 RLHF 训练提供能源并进行优化

DeepSpeed-Chat 流程的前两步与大型模型的惯例微调类似，得益于基于 ZeRO 的内存治理优化和 DeepSpeed 训练中的并行策略灵便组合，实现了规模和速度的晋升。然而，流程的第三步在性能方面是最具挑战性的局部。每次迭代都须要高效解决两个阶段：a) 生成答复的推理阶段，为训练提供输出；b) 更新 actor 和 reward 模型权重的训练阶段，以及它们之间的交互和调度。这引入了两个次要艰难：（1）内存老本，因为在第三阶段的整个过程中须要运行多个 SFT 和 RW 模型；（2）生成答复阶段的速度较慢，如果没有正确减速，将显著拖慢整个第三阶段。此外，咱们在第三阶段中增加的两个重要可选性能，包含指数挪动均匀（EMA）收集和混合训练，将产生额定的内存和训练老本。

为了应答这些挑战，咱们将 DeepSpeed 训练和推理的零碎性能整合为一个对立的基础设施，称为混合引擎（Hybrid Engine）。它利用原始 DeepSpeed 引擎进行高速训练模式，同时轻松利用 DeepSpeed 推理引擎进行生成 / 评估模式，为第三阶段的 RLHF 训练提供了一个显著更快的训练零碎。如图 2 所示，DeepSpeed 训练和推理引擎之间的过渡是无缝的：通过为 actor 模型启用典型的 eval 和 train 模式，当运行推理和训练流程时，DeepSpeed 抉择其不同的优化来运行模型更快并进步整个零碎吞吐量。

Figure 2. 设计图解：DeepSpeed Hybrid Engine，用于减速 RLHF 流程中最耗时的局部。

在 RLHF 训练的教训生成阶段的推理执行过程中，DeepSpeed 混合引擎应用轻量级内存管理系统来解决 KV 缓存和两头后果，同时应用高度优化的推理 CUDA 核和张量并行计算。与现有解决方案相比，DeepSpeed-HE 显著进步了吞吐量（每秒 token 数）。

在训练执行过程中，混合引擎应用了多种内存优化技术，如 DeepSpeed 的 ZeRO 系列技术和当初风行的 LoRA 办法。这些技术在混合引擎中能够彼此兼容，并能够组合在一起以提供最高训练效率。

DeepSpeed-HE 能够在训练和推理之间无缝更改模型分区，以反对基于张量并行计算的推理和基于 ZeRO 的分片机制进行训练。它还会重新配置内存零碎以在此期间最大化内存可用性。DeepSpeed-HE 还通过躲避内存调配瓶颈和反对大批量大小来进一步提高性能。混合引擎集成了 DeepSpeed 训练和推理的一系列零碎技术，冲破了现有 RLHF 训练的极限，并为 RLHF 工作负载提供了无可比拟的规模和零碎效率。

5. DeepSpeed RLHF: 通过 Hybrid Engine 实现无可比拟的规模和效率

回顾

如前所述，DeepSpeed-HE 是一个将弱小的用于推理和训练的联合零碎，旨在使 DeepSpeed-RLHF 在各种硬件上实现卓越的规模和效率，使 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 所示。

与现有 RLHF 零碎的吞吐量和模型大小可扩展性比拟

与其余 RLHF 零碎（如 Colossal-AI 或由原生 PyTorch 提供反对的 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 领有在雷同工夫估算下训练更大的 actor 模型的能力，或者以十分之一的老本训练相似大小的模型的能力。

图 3. 在单个 NVIDIA A100-40G GPU 上，将 RLHF 训练的吞吐量与另外两个零碎框架在步骤 3 进行比拟。没有图标示意 OOM（内存不足）的状况

图 4. 在单个 DGX 节点上，应用 8 个 NVIDIA A100-40G GPU，对训练流程第 3 步（耗时最长的局部）的不同模型大小进行端到端训练吞吐量比拟。没有图标示意 OOM（内存不足）的状况。

这种效率的进步是 DeepSpeed-HE 利用 DeepSpeed 推理优化在 RLHF 处理过程中减速 RLHF 生成的后果。图 5 显示了 RLHF 训练迭代中 1.3B 参数模型的工夫耗费细节：大部分工夫用于生成阶段。通过利用 DeepSpeed 的高性能推理内核，DeepSpeed-HE 在这个阶段能够实现比 HuggingFace 高达 9 倍的吞吐量改良，比 Colossal-AI 高 15 倍，从而实现无可比拟的端到端效率。

图 5. DeepSpeed Chat 的混合引擎在生成阶段的优越减速：在单个 DGX 节点上应用 8 个 A100-40G GPU 训练 OPT-1.3B actor 模型 + OPT-350M reward 模型的工夫 / 序列合成。

无效吞吐量和可扩展性剖析

(I) 无效吞吐量剖析。 在 RLHF 训练的第 3 阶段，DeepSpeed-HE 的无效吞吐量取决于它在生成和 RL 训练阶段所实现的吞吐量。在咱们的 RLHF（详见 benchmarking setting）中，生成阶段占总计算的约 20%，而 RL 训练阶段占残余的 80%。然而，只管比例较小，前者可能会占用大部分的端到端工夫，因为它须要为每个生成的字符运行一次 actor 模型，使其受到内存带宽限度，难以实现高吞吐量。相比之下，RL 训练阶段是计算密集型的，仅需运行参考 actor 模型进行几次前向和后向传递，每个样本都有来自提醒和生成的全副 512 个字符，能够实现良好的吞吐量。

图 6. 在最大效率的状况下，DeepSpeed-HE 针对不同模型大小的 RLHF 生成、训练和无效吞吐量。

为了最大化无效吞吐量，DeepSpeed-HE 对两个阶段进行了优化。首先，它应用尽可能大的批量大小以在两个阶段上取得更高的效率。其次，在生成阶段，它利用高性能 CUDA 内核在模型在单个 GPU 上最大化 GPU 内存带宽利用率，并在其余状况下利用张量并行（Tensor Parallelism, 简写作 TP）进行计算。DeepSpeed-HE 进一步在生成阶段应用 TP 而不是 ZeRO 以缩小 GPU 之间的通信并放弃高 GPU 内存带宽利用率。

图 6 显示了 DeepSpeed-HE 在 1.3B 到 175B 的模型大小范畴内能够实现的最佳无效吞吐量（以 TFlops/GPU 示意）。它还别离显示了在生成和训练阶段实现的吞吐量。DeepSpeed-HE 对 6.7B-66B 范畴内的模型最为高效。超出这个范畴到 175B 时，因为内存无限，无奈反对更大的批量大小，吞吐量降落，但仍比小型 1.3B 模型的效率高 1.2 倍。当咱们将这些微小的模型扩大到更多具备更多内存的 GPU 时，这些模型的每个 GPU 吞吐量可能会进一步提高。

此外，咱们想指出，如图 2 所示，咱们零碎的有效性能比现有零碎高 19 倍，这表明它们的运行速度低于峰值的 5%。这阐明了优化 RLHF 工作负载的挑战以及咱们的零碎在面对挑战时的有效性。

图 7. 在不同数量的 DGX (A100-40/80G GPU) 节点上，进行 13B（左）和 66B（右）actor 模型 和 350M reward 模型的可扩展性训练。

(II) 可扩展性剖析。 不同模型大小的最佳无效吞吐量取决于不同的 GPU 数量。局部起因是因为一些较大的模型大小须要更多的内存来运行。基于此，咱们接下来探讨 DeepSpeed-HE 的可扩展性个性。

图 7 显示 DeepSeed-RLHF 在多达 64 个 GPU 的集群 上实现了良好的整体扩大。然而，如果咱们仔细观察，能够发现 DeepSpeed-RLHF 训练在小规模时实现了超线性扩大，随后在较大规模时实现了靠近线性或次线性扩大。这是因为内存可用性和最大全局批量大小之间的相互作用。

DeepSpeed-HE 的核心技术基于 ZeRO，用于训练过程中将模型状态宰割到每个 GPU 上。这意味着随着 GPU 数量的减少，每个 GPU 的内存耗费会缩小，使得 DeepSpeed-HE 可能在每个 GPU 上反对更大的批量，从而实现超线性扩大。然而，在大规模状况下，只管可用内存继续减少，但最大全局批量大小依然限度了每个 GPU 的批量大小，导致靠近线性或次线性扩大。因而，在给定的最大全局批量大小（例如，咱们设置为 1024 个句子，每个句子长度为 512）下，DeepSpeed-HE 在超线性和次线性可扩展性之间实现了最佳的吞吐量和老本效益。具体的平衡点次要取决于每个 GPU 上可运行的最大批量大小，而这又受到可用内存和全局批量大小的函数所决定。

本文参加了 SegmentFault 思否写作挑战赛，欢送正在浏览的你也退出。