加速 RLHF 训练的自适应放置与并行框架
机器学习
2024-10-15 v3 人工智能
计算与语言
摘要
最近，类似 ChatGPT 或 InstructGPT 的大型语言模型（LLM）在人工智能领域产生了重大影响。许多工作试图复现复杂的 InstructGPT 训练流程，即基于人类反馈的强化学习（RLHF）。然而，主流的分布式 RLHF 训练方法通常采用固定的模型放置策略，称为共置策略。该策略将 RLHF 中涉及的所有四个相互依赖的模型视为单一实体，将它们分布在所有设备上，并应用为单一模型设计的并行技术，而不考虑每个模型固有的工作负载异构性。因此，该策略加剧了 RLHF 训练中的生成瓶颈，并降低了整体训练效率。为了解决这些问题，我们提出了一个灵活的模型放置框架，提供两种通用且敏捷的模型放置策略。交错策略通过精心编排将无依赖的模型放置在独占设备上，有助于减少 RLHF 训练的内存冗余和通信成本。另一方面，分离策略通过使用额外的影子模型分离 RLHF 流程的训练和推理运行时，提高了模型训练的吞吐量。此外，我们的框架提供了一个简单的用户界面和指南，以便在各种训练场景中轻松灵活地配置这些策略。我们的实验表明，与当前最先进的（SOTA）方法相比，我们的策略可以实现高达 11 倍的显著提升。结果突出了我们在加速分布式 RLHF 训练方面的有效性和适应性。
引用
@article{arxiv.2312.11819,
title = {An Adaptive Placement and Parallelism Framework for Accelerating RLHF Training},
author = {Youshao Xiao and Zhenglei Zhou and Fagui Mao and Weichang Wu and Shangchun Zhao and Lin Ju and Lei Liang and Xiaolu Zhang and Jun Zhou},
journal= {arXiv preprint arXiv:2312.11819},
year = {2024}
}