作者 | 微软 DeepSpeed 开源项目组

来源 | 开源社（中文版授权开源社首发）

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

卓越的扩展性：DeepSpeed-HE 能够支持训练拥有数千亿参数的模型，并在多节点多 GPU 系统上展现出卓越的扩展性。因此，即使是一个拥有 130 亿参数的模型，也只需 1.25 小时就能完成训练。而对于庞大的 拥有1750 亿参数的模型，使用 DeepSpeed-HE 进行训练也只需不到一天的时间。

完整的 RLHF 训练流程概述

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

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

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

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

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

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 模型的能力，或者以十分之一的成本训练类似大小的模型的能力。

图 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 的内存消耗会减少，使得 DeepSpeed-HE 能够在每个 GPU 上支持更大的批量，从而实现超线性扩展。然而，在大规模情况下，尽管可用内存持续增加，但最大全局批量大小仍然限制了每个 GPU 的批量大小，导致接近线性或次线性扩展。因此，在给定的最大全局批量大小（例如，我们设置为 1024 个句子，每个句子长度为 512）下，DeepSpeed-HE 在超线性和次线性可扩展性之间实现了最佳的吞吐量和成本效益。具体的平衡点主要取决于每个 GPU 上可运行的最大批量大小，而这又受到可用内存和全局批量大小的函数所决定。

发布：现在就尝试 DeepSpeed Chat 吧!

DeepSpeed-Chat 现已开源并向 AI 社区开放。

如果你发现这项成果对你有用或者喜欢微软的开源成果，请在 DeepSpeed（https://github.com/microsoft/DeepSpeed） 和 DeepSpeedExamples（https://github.com/microsoft/DeepSpeedExamples）上点⭐。

请访问 DeepSpeed-Chat GitHub页面以开始使用：GitHub 登陆页面（https://github.com/microsoft/DeepSpeedExamples/tree/master/applications/DeepSpeed-Chat）

研发团队将继续根据你的反馈和支持改进 DeepSpeed-Chat。下面的计划图（https://github.com/microsoft/DeepSpeedExamples/blob/master/applications/DeepSpeed-Chat/README.md#-deepspeed-chats-roadmap-）显示了当前支持的功能以及计划在未来支持的功能。

请访问官方网站（https://www.deepspeed.ai/），了解详细的博客文章、教程和有用的文档。

你还可以关注 DeepSpeed-Chat 的英文 Twitter（DeepSpeed） 和日文 Twitter（マイクロソフトDeepSpeed），了解 DeepSpeed 的最新动态。

本博客由微软DeepSpeed组官方撰写并译制（https://github.com/microsoft/DeepSpeed/tree/master/blogs/deepspeed-chat），并授权开源社首发于中文社区，如需转载请注明出处。

花粉社群VIP加油站