[论文解读] DeepSpeed-Chat: Easy, Fast and Affordable RLHF Training of ChatGPT-like Models at All Scales

Zhewei Yao, Reza Yazdani Aminabadi+17 more|arXiv (Cornell University)|Aug 2, 2023

Machine Learning and Data Classification被引用 10

一句话总结

介绍 DeepSpeed-Chat，这是一个统一的开源 RLHF 训练流水线和 Hybrid Engine，使从数十亿参数到非常大规模的 ChatGPT 类模型在单 GPU 或多节点集群上，易于、快速且具成本效益地训练。

ABSTRACT

ChatGPT-like models have revolutionized various applications in artificial intelligence, from summarization and coding to translation, matching or even surpassing human performance. However, the current landscape lacks an accessible, efficient, and cost-effective end-to-end RLHF (Reinforcement Learning with Human Feedback) training pipeline for these powerful models, particularly when training at the scale of billions of parameters. This paper introduces DeepSpeed-Chat, a novel system that democratizes RLHF training, making it accessible to the AI community. DeepSpeed-Chat offers three key capabilities: an easy-to-use training and inference experience for ChatGPT-like models, a DeepSpeed-RLHF pipeline that replicates the training pipeline from InstructGPT, and a robust DeepSpeed-RLHF system that combines various optimizations for training and inference in a unified way. The system delivers unparalleled efficiency and scalability, enabling training of models with hundreds of billions of parameters in record time and at a fraction of the cost. With this development, DeepSpeed-Chat paves the way for broader access to advanced RLHF training, even for data scientists with limited resources, thereby fostering innovation and further development in the field of AI.

研究动机与目标

• 通过提供一个可访问的端到端流水线，使 ChatGPT-类模型的 RLHF 训练民主化。
• 复制 InstructGPT 的三步 RLHF 工作流（SFT、Reward Model Fine-tuning、RLHF），通过数据抽象与混合实现。
• 提供统一而高效的系统（Hybrid Engine），优化 RLHF 负载下的训练与推理。
• 展示从数十亿到数千亿参数的模型在不同规模上的可扩展性和成本效益。）

提出的方法

• 提供一个易于使用的脚本，从经过 HuggingFace 预训练的模型开始，运行 InstructGPT- 类训练的三个 RLHF 步骤。
• 实现 DeepSpeed-RLHF 流水线，遵循 SFT、Reward Model Fine-tuning、RLHF，并具备数据抽象和混合能力。
• 开发 DeepSpeed-Hybrid Engine (DeepSpeed-HE)，结合快速推理内核、张量并行、基于 ZeRO 的内存优化，以及 LoRA，实现高效的 RLHF 训练与生成。
• 采用 EMA 检查点和 Mixture Training 以提升最终模型质量并保留预训练能力。
• 提供灵活的 API，通过可重复使用的引擎和 PPO 训练器定制 RLHF 流水线，支持研究试验。
• 在单 GPU 和多节点设置下，与 Colossal-AI 和 HuggingFace DDP 进行吞吐量和可扩展性基准对比，突出改进。

实验结果

研究问题

• RQ1 如何在不同规模上使 ChatGPT-类模型的 RLHF 训练变得更易访问、快速且具成本效益？
• RQ2 哪些系统设计与优化能够实现面向大规模 actor/reward 模型的端到端 RLHF（SFT、RM 调优、RLHF）？
• RQ3 统一的 Hybrid Engine 在生成阶段和训练阶段相较于现有框架的表现如何？
• RQ4 在从数十亿到数千亿参数的模型训练时，实际的成本、时间与可扩展性有何优势？
• RQ5 用户是否可以通过灵活的 API 自定义 RLHF 流水线，以探索新的 RLHF 策略？

主要发现

• DeepSpeed-HE 在报道的设置中实现了比现有系统快超过 15 倍的 RLHF 训练。
• 单节点 8x A100-40G 在 Azure 上分别以低于 $300 和 $600 的成本训练 OPT-13B 和 OPT-30B，耗时分别为 9 小时和 18 小时。
• 多节点 64x A100-80G 可以在 1.25 小时训练 OPT-13B、在 20 小时训练 OPT-175B，总成本约为 $5120。
• DeepSpeed-HE 使从 13B 到 175B 参数的模型在可扩展的硬件上可训练，且单 GPU 也支持 >13B 模型。
• 比较显示 DeepSpeed-HE 在多-GPU 设置下对 Colossal-AI 提供 6–19x 的加速，对 HuggingFace DDP 提供 1.4–10.5x 的加速，并且每个硬件可用的模型规模高达 7.5x 增大。