Recently, ChatGPT or InstructGPT like large language models (LLM) has made a significant impact in the AI world. Many works have attempted to reproduce the complex InstructGPT's training pipeline, namely Reinforcement Learning with Human Feedback (RLHF). However, the mainstream distributed RLHF training methods typically adopt a fixed model placement strategy, referred to as the Flattening strategy. This strategy treats all four interdependent models involved in RLHF as a single entity, distributing them across all devices and applying parallelism techniques designed for a single model, regardless of the different workloads inherent to each model. As a result, this strategy exacerbates the generation bottlenecks in the RLHF training and degrades the overall training efficiency. To address these issues, we propose an adaptive model placement framework that offers two flexible model placement strategies. The Interleaving strategy helps reduce memory redundancy and communication costs of RLHF training by placing models without dependencies on exclusive devices with careful orchestration. On the other hand, the Separation strategy improves the throughput of model training by separating the training and inference runtime of the RLHF pipeline with additional shadow models. Furthermore, our framework provides a simple user interface and allows for the agile allocation of models across devices in a fine-grained manner for various training scenarios, involving models of varying sizes and devices of different scales. Extensive experiments have demonstrated that our Interleaving and Separation strategies can achieve notable improvements up to 11X, compared to the current SOTA approaches. The results highlight the effectiveness and adaptability of our approaches in accelerating the training of distributed RLHF.
翻译:近期,ChatGPT或InstructGPT类大语言模型(LLM)在人工智能领域产生了深远影响。许多研究尝试复现复杂的InstructGPT训练流程,即基于人类反馈的强化学习(RLHF)。然而,当前主流的分布式RLHF训练方法通常采用固定的模型部署策略,即“扁平化策略”。该策略将RLHF涉及的四个相互依赖模型视为单一整体,将其分布于所有设备上,并应用针对单一模型设计的并行技术,而忽略各模型固有的工作负载差异。这导致RLHF训练中的生成瓶颈加剧,整体训练效率下降。针对上述问题,我们提出了一种自适应模型部署框架,提供两种灵活的模型部署策略:交错策略通过将无依赖关系的模型精心编排至独立设备部署,有效降低RLHF训练的内存冗余与通信开销;分离策略则通过引入额外的影子模型分离RLHF流水线的训练与推理运行时,提升模型训练吞吐量。此外,本框架提供简洁的用户接口,支持针对不同训练场景(涵盖不同规模模型与异构设备)的细粒度模型敏捷分配。大量实验表明,相较于现有最先进方法,我们的交错策略与分离策略可实现高达11倍的显著性能提升。实验结果充分验证了所提方法在加速分布式RLHF训练中的有效性与适应性。