人手一个 ChatGPT，微软 DeepSpeed Chat 震撼发布，一键 RLHF 训练千亿级大模型

人手一个 ChatGPT 的梦想，就要实现了？

刚刚，微软开源了一个可以在模型训练中加入完整 RLHF 流程的系统框架 ——DeepSpeed Chat。

也就是说，各种规模的高质量类 ChatGPT 模型，现在都唾手可得了！

一键解锁千亿级 ChatGPT，轻松省钱 15 倍

众所周知，由于 OpenAI 太不 Open，开源社区为了让更多人能用上类 ChatGPT 模型，相继推出了 LLaMa、Alpaca、Vicuna、Databricks-Dolly 等模型。

但由于缺乏一个支持端到端的 RLHF 规模化系统，目前类 ChatGPT 模型的训练仍然十分困难。而 DeepSpeed Chat 的出现，正好补全了这个「bug」。

更亮的是，DeepSpeed Chat 把成本大大地打了下来。

此前，昂贵的多 GPU 设置超出了许多研究者的能力范围，并且，即使能访问多 GPU 集群，现有的方法也无力负担数千亿参数 ChatGPT 模型的训练。

现在，只要花 1620 美元，就可以通过混合引擎 DeepSpeed-HE，在 2.1 天内训练一个 OPT-66B 模型。

而如果使用多节点、多 GPU 系统，DeepSpeed-HE 可以花 320 美元，在 1.25 小时内训练一个 OPT-13B 模型，花 5120 美元，就能在不到一天的时间内训练一个 OPT-175B 模型。

前 Meta AI 专家 Elvis 激动转发，称这是一件大事，并表示好奇 DeepSpeed Chat 和 ColossalChat 相比起来如何。

下面，我们就来看看效果如何。

经过 DeepSpeed-Chat 的训练，13 亿参数版「ChatGPT」在问答环节上的表现非常亮眼。不仅能 get 到问题的上下文关系，而且给出的答案也有模有样。

在多轮对话中，这个 13 亿参数版「ChatGPT」所展示出的性能，也完全超越了这个规模的固有印象。

一段代码，生成你的第一个 ChatGPT

当然，在体验之前，还需要把环境配置一下：

git clone https://github.com/microsoft/DeepSpeed.gitcd DeepSpeedpip install .git clone https://github.com/microsoft/DeepSpeedExamples.gitcd DeepSpeedExamples/lications/DeepSpeed-Chat/pip install -r requirements.txt

一杯咖啡，训完 13 亿参数版 ChatGPT

如果你只有大约 1-2 小时的咖啡或午餐休息时间，也可以尝试使用 DeepSpeed-Chat 训练一个「小玩具」。

团队特地准备了一个针对 1.3B 模型的训练示例，可在消费级 GPU 上进行测试。最棒的是，当你从午休回来时，一切都已准备就绪。

python train.py --actor-model facebook/opt-1.3b --reward-model facebook/opt-350m --num-gpus 1

配备 48GB 显存的消费级 NVIDIA A6000 GPU：

一个 GPU Node，半天搞定 130 亿参数

如果你只有半天的时间，以及一台服务器节点，则可以通过预训练的 OPT-13B 作为 actor 模型，OPT-350M 作为 reward 模型，来生成一个 130 亿参数的类 ChatGPT 模型：

python train.py --actor-model facebook/opt-13b --reward-model facebook/opt-350m --num-gpus 8

单 DGX 节点，搭载了 8 个 NVIDIA A100-40G GPU：

超省钱云方案，训练 660 亿参数模型

如果你可以使用多节点集群或云资源，并希望训练一个更大、更高质量的模型。那么只需基于下面这行代码，输入你想要的模型大小（如 66B）和 GPU 数量（如 64）：

python train.py --actor-model facebook/opt-66b --reward-model facebook/opt-350m --num-gpus 64

8 个 DGX 节点，每个节点配备 8 个 NVIDIA A100-80G GPU：

具体来说，针对不同规模的模型和硬件配置，DeepSpeed-RLHF 系统所需的时间和成本如下：

DeepSpeed Chat 是个啥？

DeepSpeed Chat 是一种通用系统框架，能够实现类似 ChatGPT 模型的端到端 RLHF 训练，从而帮助我们生成自己的高质量类 ChatGPT 模型。

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

1. 简化 ChatGPT 类型模型的训练和强化推理体验

开发者只需一个脚本，就能实现多个训练步骤，并且在完成后还可以利用推理 API 进行对话式交互测试。

2. DeepSpeed-RLHF 模块

DeepSpeed-RLHF 复刻了 InstructGPT 论文中的训练模式，并提供了数据抽象和混合功能，支持开发者使用多个不同来源的数据源进行训练。

3. DeepSpeed-RLHF 系统

团队将 DeepSpeed 的训练（training engine）和推理能力（inference engine) 整合成了一个统一的混合引擎（DeepSpeed Hybrid Engine or DeepSpeed-HE）中，用于 RLHF 训练。由于，DeepSpeed-HE 能够无缝地在推理和训练模式之间切换，因此可以利用来自 DeepSpeed-Inference 的各种优化。

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

高效且经济：

DeepSpeed-HE 比现有系统快 15 倍以上，使 RLHF 训练快速且经济实惠。

例如，DeepSpeed-HE 在 Azure 云上只需 9 小时即可训练一个 OPT-13B 模型，只需 18 小时即可训练一个 OPT-30B 模型。这两种训练分别花费不到 300 美元和 600 美元。

卓越的扩展性：

DeepSpeed-HE 能够支持训练拥有数千亿参数的模型，并在多节点多 GPU 系统上展现出卓越的扩展性。

因此，即使是一个拥有 130 亿参数的模型，也只需 1.25 小时就能完成训练。而对于拥有 1750 亿参数的模型，使用 DeepSpeed-HE 进行训练也只需不到一天的时间。

实现 RLHF 训练的普及化：

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

完整的 RLHF 训练流程

为了提供无缝的训练体验，研究者遵循 InstructGPT，并在 DeepSpeed-Chat 中包含了一个完整的端到端训练流程。

DeepSpeed-Chat 的 RLHF 训练流程图示，包含了一些可选择的功能

流程包括三个主要步骤：

第 1 步：

监督微调 (SFT)，使用精选的人类回答来微调预训练的语言模型，以应对各种查询。

第 2 步：

奖励模型微调，用一个包含人类对同一查询的多个答案打分的数据集，来训练一个独立的（通常比 SFT 小）奖励模型（RW）。

第 3 步：

RLHF 训练，在这一步，SFT 模型通过使用近似策略优化（PPO）算法，从 RW 模型的奖励反馈进一步微调。

在步骤 3 中，研究者还提供了两个附加功能，来帮助提高模型质量：

- 指数移动平均线（EMA）的收集，可以选择一个基于 EMA 的检查点，进行最终评估。

- 混合训练，将预训练目标（即下一个词预测）与 PPO 目标混合，以防止在公共基准（如 SQuAD2.0）上的性能回归。

EMA 和混合训练这两个训练特征，常常被其他的开源框架所忽略，因为它们并不会妨碍训练的进行。

然而，根据 InstructGPT，EMA 检查点往往比传统的最终训练模型提供更好的响应质量，而混合训练可以帮助模型保持训练前的基准解决能力。

因此，研究者为用户提供了这些功能，让他们可以充分获得 InstructGPT 中描述的训练经验。

而除了与 InstructGPT 论文高度一致外，研究者还提供了功能，让开发者使用多种数据资源，训练自己的 RLHF 模型：

数据抽象和混合能力：

DeepSpeed-Chat 配备了（1）抽象数据集层，以统一不同数据集的格式；以及（2）数据拆分 / 混合功能，从而使多个数据集被适当地混合，然后在 3 个训练阶段进行分割。

DeepSpeed 混合引擎

指示引导的 RLHF 管道的第 1 步和第 2 步，类似于大模型的常规微调，它们由基于 ZeRO 的优化和 DeepSpeed 训练中灵活的并行策略组合，来实现规模和速度。

而管道的第 3 步是在性能影响方面最复杂的部分。

每个迭代都需要有效地处理两个阶段：a）推理阶段，用于 token / 经验的生成，产生训练的输入；b）训练阶段，更新 actor 和 reward 模型的权重，以及它们之间的互动和调度。

它引入了两个主要的困难：（1）内存成本，因为在整个第三阶段需要运行多个 SFT 和 RW 模型；（2）生成回答阶段的速度较慢，如果没有正确加速，将显著拖慢整个第三阶段。

此外，研究者在第三阶段增加的两个重要功能 —— 指数移动平均（EMA）收集和混合训练，将产生额外的内存和训练成本。

为了应对这些挑战，研究者将 DeepSpeed 训练和推理的全部系统能力组成了一个统一的基础设施，即混合引擎（Hybrid Engine）。

它利用原有的 DeepSpeed 引擎进行快速训练模式，同时毫不费力地应用 DeepSpeed 推理引擎进行生成 / 评估模式，为第三阶段的 RLHF 训练提供了一个更快的训练系统。

如下图所示，DeepSpeed 训练和推理引擎之间的过渡是无缝的：通过为 actor 模型启用典型的 eval 和 train 模式，在运行推理和训练流程时，DeepSpeed 选择了不同的优化，以更快地运行模型，并提高整个系统的吞吐量。 

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

在训练期间，混合引擎启用了内存优化技术，如 DeepSpeed 的 ZeRO 系列技术和低阶自适应（LoRA）。

而研究者设计和实现这些系统优化的方式是，让它们彼此兼容，并可以组合在一起，在统一的混合引擎下提供最高的训练效率。

混合引擎可以在训练和推理中无缝地改变模型分区，以支持基于张量并行的推理，和基于 ZeRO 的训练分片机制。

它还可以重新配置内存系统，以便在每一种模式中最大限度地提高内存可用性。

这就避免了内存分配瓶颈，能够支持大的 batch size，让性能大大提升。

总之，混合引擎推动了现代 RLHF 训练的边界，为 RLHF 工作负载提供了无与伦比的规模和系统效率。

效果评估

与 Colossal-AI 或 HuggingFace-DDP 等现有系统相比，DeepSpeed-Chat 具有超过一个数量级的吞吐量，能够在相同的延迟预算下训练更大的演员模型或以更低的成本训练相似大小的模型。

例如，在单个 GPU 上，DeepSpeed 使 RLHF 训练的吞吐量提高了 10 倍以上。虽然 CAI-Coati 和 HF-DDP 都可以运行 1.3B 的模型，但 DeepSpeed 可以在相同的硬件上运行 6.5B 模型，直接高出 5 倍。

在单个节点的多个 GPU 上，DeepSpeed-Chat 在系统吞吐量方面比 CAI-Coati 提速 6-19 倍，HF-DDP 提速 1.4-10.5 倍。

团队表示，DeepSpeed-Chat 能够获得如此优异的结果，关键原因之一便是混合引擎在生成阶段提供的加速。

参考资料：

• https://github.com/microsoft/DeepSpeed

本文来自微信公众号：新智元 （ID：AI_era）

广告声明：文内含有的对外跳转链接（包括不限于超链接、二维码、口令等形式），用于传递更多信息，节省甄选时间，结果仅供参考，IT之家所有文章均包含本声明。