一键式 RLHF 训练 DeepSpeed Chat（一）：理论篇

背景

近日，类ChatGPT模型引发了人工智能（AI）领域的一场风潮。 这场风潮对数字世界产生了革命性影响。类ChatGPT模型具有惊人的泛用性，能够执行归纳、编程、翻译等任务，其结果与人类专家相当甚至更优。为了使ChatGPT等模型的训练和部署更轻松，AI 开源社区进行了各种尝试（例如：ChatLLaMa、Alpaca、Vicuna等）。

尽管开源社区付出了巨大的努力，目前仍缺乏一个支持端到端的基于人工反馈机制的强化学习（RLHF）的规模化系统，这使得训练强大的类ChatGPT模型十分困难。例如，使用现有的开源系统训练一个具有 67 亿参数的类ChatGPT模型通常需要昂贵的多卡至多节点的 GPU 集群，但这些资源对大多数数据科学家或研究者而言难以获取。同时，即使有了这样的计算资源，现有的开源系统的训练效率通常还不到这些机器所能达到的最大效率的5%。

简而言之，即使有了昂贵的多GPU集群，现有解决方案也无法轻松、快速、经济的训练具有数千亿参数的最先进的类ChatGPT模型。

ChatGPT模型的训练是基于InstructGPT论文中的RLHF方式。这与常见的大语言模型的预训练和微调截然不同。这使得现有深度学习系统在训练类ChatGPT模型时存在种种局限。

为了让ChatGPT类型的模型更容易被普通数据科学家和研究者使用，并使RLHF训练真正普及到AI社区，DeepSpeed-Chat应运而生。

DeepSpeed Chat 核心功能

DeepSpeed Chat具有以下三大核心功能：

• 简化 ChatGPT 类型模型的训练和强化推理体验：只需一个脚本即可实现多个训练步骤，包括使用 Huggingface 预训练的模型、使用 DeepSpeed-RLHF 系统运行 InstructGPT 训练的所有三个步骤、甚至生成你自己的类ChatGPT模型。此外，还提供了一个易于使用的推理API，用于用户在模型训练后进行对话式交互测试。
• DeepSpeed RLHF 模块：DeepSpeed RLHF 复刻了 InstructGPT 论文中的训练模式，并确保包括监督微调（SFT），奖励模型微调和基于人类反馈的强化学习（RLHF）在内的三个步骤并与其一一对应。此外，还提供了数据抽象和混合功能，以支持用户使用多个不同来源的数据源进行训练。
• DeepSpeed RLHF 系统：将 DeepSpeed 的训练（training engine）和推理能力（inference engine) 整合到一个统一的混合引擎（DeepSpeed Hybrid Engine / DeepSpeed-HE）中用于 RLHF 训练。DeepSpeed-HE 能够在 RLHF 中无缝地在推理和训练模式之间切换，使其能够利用来自 DeepSpeed-Inference 的各种优化，如张量并行计算和高性能CUDA算子进行语言生成，同时对训练部分还能从 ZeRO 和 LoRA-based 内存优化策略中受益。DeepSpeed-HE 还能够自动在 RLHF 的不同阶段进行智能的内存管理和数据缓存。

DeepSpeed RLHF 系统的优势

DeepSpeed-RLHF 系统在大规模训练中具有无与伦比的效率使复杂的 RLHF 训练变得快速、经济并且易于大规模推广。

高效性和经济性：DeepSpeed-HE 比现有系统快 15 倍以上，使 RLHF 训练快速且经济实惠。例如，DeepSpeed-HE 在 Azure 云上只需 9 小时即可训练一个 OPT-13B模型，只需 18 小时即可训练一个 OPT-30B模型。这两种训练分别花费不到 300 美元和 600 美元。 下表展示了单节点 8 卡A100的训练时长及预估的 Azure 费用。

极好的扩展性：DeepSpeed-HE 能够支持训练拥有数千亿参数的模型，并在多节点多 GPU 系统上展现出极好的扩展性。因此，即使是一个拥有 130 亿参数的模型，也只需 1.25 小时就能完成训练。而对于庞大的 拥有1750 亿参数的模型，使用 DeepSpeed-HE 进行训练也只需不到一天的时间。下表展示了多节点 64 卡 A100-80GB 的训练时长及预估的 Azure 费用。

注意：

上述两个表格中的数据均针对 RLHF 训练的第三阶段，基于实际数据集和 DeepSpeed-RLHF 训练吞吐量的测试。该训练在总共 1.35 亿（135M）个token上进行一个 epoch 的训练。总共有 6750 万个查询（query）token（131.9k 个 query，每个序列长度为 256）和 6750 万个生成/回答字符（131.9k 个 answer，每个序列长度为 256），每步的最大全局字符批量大小约为 500 万个字符（1024 个query-answer对）。

实现 RLHF 训练的普及化：仅凭单个 GPU，DeepSpeed-HE 就能支持训练超过 130 亿参数的模型。这使得那些无法使用多 GPU 系统的数据科学家和研究者不仅能够轻松创建轻量级的 RLHF 模型，还能创建大型且功能强大的模型以应对不同的使用场景。下表展示了 DeepSpeed-HE 在不同的 GPU 单卡上支持的最大模型。

DeepSpeed Chat 完整的 RLHF 训练流程

为了实现无缝的训练体验，遵循 InstructGPT 论文的方法，DeepSpeed-Chat 中整合了一个端到端的训练流程。下图展示了 DeepSpeed-Chat 的 RLHF 训练流程，并包含了一些可选择的功能。

流程包括三个主要阶段：

• 第一阶段：监督微调（SFT）  —— 使用精选的人类回答来微调预训练的语言模型以应对各种查询（query）；
• 第二阶段：奖励模型微调 —— 使用一个包含人类对同一查询（query）的多个答案打分的数据集来训练一个独立的（通常比 SFT 小的）奖励模型（RW）；
• 第三阶段：RLHF 训练 —— 利用 Proximal Policy Optimization（PPO）算法，根据 RW 模型的奖励反馈进一步微调 SFT 模型。

在第三阶段中，还提供了两个额外的功能，以帮助提高模型质量：

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

这两个训练功能(EMA 和混合训练)常常被其他的开源框架所忽略，因为它们并不会妨碍训练的进行。然而，根据 InstructGPT 描述，EMA 通常比传统的最终训练模型提供更好的响应质量，而混合训练可以帮助模型保持预训练基准解决能力。因此，DeepSpeed Chat为用户提供这些功能，以便充分获得 InstructGPT 中描述的训练体验，并争取更高的模型质量。

除了与 InstructGPT 论文高度一致外，DeepSpeed Chat还提供了一项方便的功能，以支持研究人员和从业者使用多个数据资源训练他们自己的 RLHF 模型：

• 数据抽象和混合能力： DeepSpeed-Chat 能够使用多个不同来源的数据集训练模型以获得更好的模型质量。它配备了（1）一个抽象数据集层，以统一不同数据集的格式；以及（2）数据拆分/混合功能，以便多个数据集在 3 个训练阶段中被适当地混合然后拆分。

DeepSpeed 混合引擎（Hybrid Engine）概述

DeepSpeed-Chat流程的前两步与大型模型的常规微调相似，得益于基于ZeRO的内存管理优化和DeepSpeed训练中的并行策略灵活组合，实现了规模和速度的提升。

然而，流程的第三阶段在性能方面是最具挑战性的部分。每次迭代都需要高效处理两个阶段：

1. 生成回答的推理阶段，为训练提供输入；
2. 更新 actor 和 reward 模型权重的训练阶段，以及它们之间的交互和调度。

这引入了两个主要困难：

1. 内存成本，因为在第三阶段的整个过程中需要运行多个SFT和RW模型；
2. 生成回答阶段的速度较慢，如果没有正确加速，将显著拖慢整个第三阶段。

此外，DeepSpeed-Chat在第三阶段中添加的两个重要的可选功能，包括指数移动平均（EMA）收集和混合训练，将产生额外的内存和训练成本。

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

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

在训练执行过程中，混合引擎使用了多种内存优化技术，如DeepSpeed的ZeRO系列技术和现在流行的LoRA方法。这些技术在混合引擎中可以彼此兼容，并可以组合在一起以保障最高训练效率。

DeepSpeed-HE可以在训练和推理之间无缝更改模型分区，以支持基于张量并行计算的推理和基于ZeRO的分片机制进行训练。它还会重新配置内存系统以在此期间最大化内存可用性。DeepSpeed-HE还通过规避内存分配瓶颈和支持大批量大小（batch size）来进一步提高性能。

混合引擎集成了DeepSpeed训练和推理的一系列系统技术，突破了现有RLHF训练的极限，并为RLHF工作负载提供了无与伦比的规模和系统效率。

DeepSpeed RLHF 性能分析

模型规模与训练成本分析

DeepSpeed RLHF具备更大的模型规模和低训练成本优势。下表针对不同actor模型大小的端到端 RLHF 训练（第 3 阶段）和固定 350M critical 模型，在硬件上运行的范围从单个消费级 GPU (NVIDIA A6000) 到更强大的云环境 (64x A100-80GPU)全面展示了 DeepSpeed-RLHF 系统支持的规模和端到端训练时间。它还展示了在 Azure Cloud 中训练模型的最具成本效益的方法以及相关成本。

DeepSpeed RLHF 与现有 RLHF 系统的吞吐量和模型大小可扩展性比较

DeepSpeed-HE 是一个强大的用于推理和训练的结合系统，旨在使 DeepSpeed-RLHF 在各种硬件上实现更大的规模和更高的效率，使 RLHF 训练快速、经济并且易于 AI 社区使用。

与其他 RLHF 系统（如 Colossal-AI 或由原生 PyTorch 提供支持的 HuggingFace）相比，DeepSpeed-RLHF 在系统性能和模型可扩展性方面表现出色：

吞吐量

就吞吐量而言，DeepSpeed 在单个 GPU 上的 RLHF 训练中实现了 10 倍以上的改进。在多 GPU 设置中，它比 Colossal-AI 快 6 - 19 倍，比 HuggingFace DDP 快 1.4 - 10.5 倍。

下图展示了在单个 NVIDIA A100-40G GPU 上，将 RLHF 训练的吞吐量与另外两个系统框架在第三阶段进行比较。其中，没有图标表示 OOM（内存不足）的情况。

下图展示了在单个 DGX 节点上，使用 8 卡 NVIDIA A100-40G GPU，对训练流程第三阶段（耗时最长的部分）的不同模型大小进行端到端训练的吞吐量比较。其中，没有图标表示 OOM（内存不足）的情况。

模型可扩展性

就模型可扩展性而言，Colossal-AI 可以在单个 GPU 上运行最大 1.3B 的模型，在单个 A100 40G 节点上运行 6.7B 的模型，而 DeepSpeed-HE 可以在相同的硬件上分别运行 6.5B 和 50B 的模型，实现高达 7.5 倍的提升。

因此，凭借超过一个数量级的更高吞吐量，与现有的 RLHF 系统（如 Colossal-AI 或 HuggingFace DDP）相比，DeepSpeed-HE 拥有在相同时间预算下训练更大的 actor 模型的能力，或者以十分之一的成本训练类似大小的模型的能力。

这种效率的提高是 DeepSpeed-HE 利用 DeepSpeed 推理优化在 RLHF 处理过程中加速 RLHF 生成的结果。下图（DeepSpeed Chat 的混合引擎在生成阶段的优越加速：在单个 DGX 节点上使用 8 卡 A100-40G GPU 训练 OPT-1.3B actor 模型 + OPT-350M reward 模型的时间/序列分解）显示了 RLHF 训练迭代中 1.3B 参数模型的时间消耗细节：大部分时间用于生成阶段。通过利用 DeepSpeed 的高性能推理内核，DeepSpeed-HE 在这个阶段可以实现比 HuggingFace 高达 9 倍的吞吐量改进，比 Colossal-AI 高 15 倍，从而实现无与伦比的端到端效率。

DeepSpeed RLHF 有效吞吐量和可扩展性分析

有效吞吐量分析

在 RLHF 训练的第三阶段，DeepSpeed-HE 的有效吞吐量取决于它在生成和 RL 训练阶段所实现的吞吐量。在DeepSpeedChat的 RLHF 中，生成阶段占总计算的约 20%，而 RL 训练阶段占剩余的 80%。尽管比例较小，前者可能会占用大部分的端到端时间，因为它需要为每个生成的字符运行一次 actor 模型，使其受到内存带宽限制，难以实现高吞吐量。相比之下，RL 训练阶段是计算密集型的，仅需运行参考（reference） actor 模型进行几次前向和后向传递，每个样本都有来自提示和生成的全部 512 个字符，可以实现良好的吞吐量。

为了最大化有效吞吐量，DeepSpeed-HE 对两个阶段进行了优化。

• 首先，它使用尽可能大的批量大小以在两个阶段上获得更高的效率。
• 其次，在生成阶段，它利用高性能CUDA内核让模型在单个 GPU 上最大化 GPU 内存带宽利用率，并在其他情况下利用张量并行（Tensor Parallelism, 简写作TP）进行计算。DeepSpeed-HE进一步在生成阶段使用 TP 而不是 ZeRO 以减少 GPU 之间的通信并保持高 GPU 内存带宽利用率。

下图显示了 DeepSpeed-HE 在 1.3B 到 175B 的模型大小范围内可以实现的最佳有效吞吐量（以 TFlops/GPU 表示）。它还分别显示了在生成和训练阶段实现的吞吐量。DeepSpeed-HE 对 6.7B-66B 范围内的模型最为高效。超出这个范围到 175B 时，由于内存有限，无法支持更大的批量大小，吞吐量下降，但仍比小型 1.3B 模型的效率高 1.2 倍。当我们将这些巨大的模型扩展到更多具有更多内存的 GPU 时，这些模型的每个 GPU 吞吐量可能会进一步提高。

可扩展性分析

不同模型大小的最佳有效吞吐量取决于不同的 GPU 数量。部分原因是因为一些较大的模型大小需要更多的内存来运行。基于此，接下来讨论 DeepSpeed-HE 的可扩展性特性。

下图（在不同数量的DGX (A100-40/80G GPU) 节点上，进行13B（左）和66B（右）actor 模型 和 350M reward 模型的可扩展性训练）显示 DeepSeed-RLHF 在多达 64 个 GPU的集群上实现了良好的整体扩展。然而，如果我们仔细观察，可以发现 DeepSpeed-RLHF 在小规模训练时实现了超线性扩展，随后在较大规模时实现了接近线性或次线性扩展。这是由于内存可用性和最大全局批量大小之间的相互作用。

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

支持的模型

目前，DeepSpeedChat支持以下模型。

当前DeepSpeedChat所有性能和精度测试均使用 OPT 模型系列进行的。对于其他模型，请查看training_scripts文件夹以了解如何更改其他模型。

DeepSpeedChat 未来规划

• 系统性地支持LLaMA微调
• 支持ZeRO-Offload (CPU/NVMe)
• 支持更广泛的 RL 算法/范式
• Auto-tuning of system optimizations
• DeepSpeed RLHF 抽象的泛化

参考文档

• DeepSpeed Chat: 一键式RLHF训练，让你的类ChatGPT千亿大模型提速省钱15倍
• DeepSpeed-Chat: Easy, Fast and Affordable RLHF Training of ChatGPT-like Models at All Scales