有效吞吐量和可扩展性分析
(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-)显示了当前支持的功能以及计划在未来支持的功能。
DeepSpeed-Chat 是更大的 DeepSpeed 生态系统的一部分,包括众多深度学习系统和建模技术。要了解更多信息,
- 请访问我们的网站(https://www.deepspeed.ai/),了解详细的博客文章、教程和有用的文档。
- 你还可以关注我们的英文 Twitter(DeepSpeed) 和日文 Twitter(マイクロソフト DeepSpeed),了解 DeepSpeed 的最新动态。
DeepSpeed 欢迎你的贡献!我们鼓励你在 DeepSpeed GitHub 页面(https://github.com/microsoft/DeepSpeed/)报告问题、贡献 PR 并参与讨论。请参阅我们的贡献指南(https://github.com/microsoft/DeepSpeed/blob/master/CONTRIBUTING.md)了解更多详情。我们愿意与大学、研究实验室、公司等进行合作,共同开展深度学习研究,将 DeepSpeed 应用于赋能现实世界的 AI 模型和应用等。对于此类需求(以及其他不适合在 GitHub 上提出的需求),请直接发送电子邮件至 deepspeed-info@microsoft.com。