ReaLHF: Optimized RLHF Training for Large Language Models through Parameter Reallocation (2406.14088v1)

Abstract: Reinforcement Learning from Human Feedback (RLHF) stands as a pivotal technique in empowering LLM applications. Since RLHF involves diverse computational workloads and intricate dependencies among multiple LLMs, directly adopting parallelization techniques from supervised training can result in sub-optimal performance. To overcome this limitation, we propose a novel approach named parameter ReaLlocation, which dynamically redistributes LLM parameters in the cluster and adapts parallelization strategies during training. Building upon this idea, we introduce ReaLHF, a pioneering system capable of automatically discovering and running efficient execution plans for RLHF training given the desired algorithmic and hardware configurations. ReaLHF formulates the execution plan for RLHF as an augmented dataflow graph. Based on this formulation, ReaLHF employs a tailored search algorithm with a lightweight cost estimator to discover an efficient execution plan. Subsequently, the runtime engine deploys the selected plan by effectively parallelizing computations and redistributing parameters. We evaluate ReaLHF on the LLaMA-2 models with up to $4\times70$ billion parameters and 128 GPUs. The experiment results showcase ReaLHF's substantial speedups of $2.0-10.6\times$ compared to baselines. Furthermore, the execution plans generated by ReaLHF exhibit an average of $26\%$ performance improvement over heuristic approaches based on Megatron-LM. The source code of ReaLHF is publicly available at https://github.com/openpsi-project/ReaLHF .

Optimizing RLHF Training for LLMs through Parameter Reallocation

This paper introduces ReaLHF, a system designed to enhance the efficiency of Reinforcement Learning from Human Feedback (RLHF) training for LLMs. The paper presents a novel approach termed parameter reallocation, enabling dynamic redistribution of LLM parameters across a GPU cluster to optimize computational workloads and address the intricate dependencies inherent in RLHF settings.

Context and Motivation

LLMs, such as GPT-3 and ChatGPT, rely heavily on extensive hardware resources due to their vast parameter sizes, driving the necessity for multiparallelization strategies to distribute computations effectively across GPUs. While traditional parallelization approaches, including data, tensor-model, and pipeline-model parallelism, are well-explored in the context of supervised training, their direct application to RLHF remains sub-par due to RLHF's distinct infrastructure requirements and multi-model dependencies.

Existing RLHF training systems often suffer from over-parallelization, leading to inefficiencies exemplified by synchronization and communication overheads in GPU clusters, or under-utilization, due to dependencies that prevent optimal GPU usage. This paper posits that parameter reallocation—dynamically adjusting the distribution of LLM parameters across devices during training—can efficiently address such bottlenecks by enabling tailored parallelization strategies for each function call type within RLHF.

Methodology

The central innovation of ReaLHF lies in its ability to automatically discover and execute efficient execution plans. It models the RLHF workflow as an augmented dataflow graph, transforming parameter reallocation and LLM execution into a systematic optimization problem.

1. Execution Plan Formulation: Each RLHF function call is assigned a device mesh and a specific 3D parallelization strategy. Execution plans are represented as augmented dataflow graphs where computations are mapped to optimal device and parallel configurations.
2. MCMC-Based Search: The exploration of execution plans leverages Markov Chain Monte Carlo (MCMC) sampling to navigate a vast combinatorial space efficiently. This method identifies cost-effective execution plans based on predicted time and memory costs while conforming to device memory constraints.
3. Runtime Execution: The chosen execution plan is operationalized on the ReaLHF system, utilizing a master-worker model to manage the dynamic redistribution of parameters across GPUs, optimizing data transfers, and ensuring efficient parallel execution.

Performance and Implications

Experimentation with LLaMA-2 models demonstrated significant speed-ups—ranging from 2.0 to 10.6 times—compared to existing systems, underscoring the efficacy of ReaLHF's innovative parameter reallocation technique. These results emphasize the capability of ReaLHF to reduce communication costs and maximize GPU utilization by exploiting concurrent execution across disjoint device subsets.

ReaLHF exhibits distinct advantages over baseline systems by dynamically adapting to varying computational patterns inherent in RLHF, such as generation and inference diversity, without the need for manually configured resource allocation and parallel strategies.

Future Perspectives

ReaLHF sets a precedent for future LLM training system designs by illustrating the potential of parameter reallocation and automated execution planning in complex RLHF workflows. The framework provides a foundation for subsequent research endeavors focusing on optimizing model training pipelines, particularly in scenarios where multi-model dependencies complicate resource management.

While promising, ReaLHF's approaches are primarily tuned for decoder-only transformer architectures and fixed workflows, leaving open research opportunities in expanding the adaptability to broader model types and dynamic dataflow configurations. Furthermore, its implementation suggests unexplored avenues in integrating emerging optimizations for single-function calls like memory-efficient attention mechanisms with ReaLHF's comprehensive execution framework.

ReaLHF exemplifies the ongoing evolution of AI system architecture, emphasizing the critical synergy between algorithmic innovation and hardware awareness to sustain the advancement of LLM applications.