Recently, ChatGPT or InstructGPT like large language models (LLM) has made a significant impact in the AI world. These models are incredibly versatile, capable of performing language tasks on par or even exceeding the capabilities of human experts. Many works have attempted to reproduce the complex InstructGPT's RLHF (Reinforcement Learning with Human Feedback) training pipeline. 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 models involved in RLHF as a single entity and places them on all devices, regardless of their differences. Unfortunately, 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. These strategies allow for the agile allocation of models across devices in a fine-grained manner. The Interleaving strategy helps reduce memory redundancy and communication costs during RLHF training. On the other hand, the Separation strategy improves the throughput of model training by separating the training and generation stages of the RLHF pipeline. Notably, this framework seamlessly integrates with other mainstream techniques for acceleration and enables automatic hyperparameter search. Extensive experiments have demonstrated that our Interleaving and Separation strategies can achieve notable improvements up to 11x, compared to the current state-of-the-art (SOTA) approaches. These experiments encompassed a wide range of training scenarios, involving models of varying sizes and devices of different scales. The results highlight the effectiveness and superiority of our approaches in accelerating the training of distributed RLHF.
翻译:近期,ChatGPT或InstructGPT类大语言模型(LLM)在人工智能领域产生了重大影响。这些模型功能极为强大,能够执行与人类专家相当甚至超越人类专家的语言任务。许多研究试图复现InstructGPT复杂的RLHF(基于人类反馈的强化学习)训练流程。然而,主流分布式RLHF训练方法通常采用固定模型放置策略,即"扁平化"策略。该策略将RLHF中涉及的四个模型视为统一整体,不加区分地放置于所有设备上。不幸的是,该策略加剧了RLHF训练中的生成瓶颈,降低了整体训练效率。为解决这些问题,我们提出一种自适应模型放置框架,提供两种灵活的模型放置策略,能以细粒度方式敏捷地在设备间分配模型。"交错"策略有助于减少RLHF训练中的内存冗余和通信开销;"分离"策略则通过分离RLHF流程中的训练与生成阶段来提升模型训练吞吐量。特别地,该框架可无缝集成其他主流的加速技术,并支持超参数自动搜索。大量实验表明,与当前最先进方法相比,我们的"交错"与"分离"策略可实现高达11倍的显著加速效果。实验涵盖了多种训练场景,涉及不同规模的模型和不同规模的设备集群。结果充分证明了我们方法在加速分布式RLHF训练中的有效性和优越性。