ReaL: Efficient RLHF Training of Large Language Models with Parameter Reallocation

ReaL: Efficient RLHF Training of Large Language Models with Parameter Reallocation

Speakers: Zhiyu Mei, Chinghua University & Ant Research; Wei Fu; Kaiwei Li; Guangju Wang; Huanchen Zhang; Yi Wu

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

The talk introduces ReaL (Reinforcement Learning with Parameter Reallocation), a novel system designed to significantly enhance the efficiency of Reinforcement Learning from Human Feedback (RLHF) for large language models (LLMs). Presented by Zhiyu Mei from Chinghua University and Ant Research, this work addresses the complex computational demands posed by RLHF, which involves multiple interacting models and distinct computational stages, unlike traditional supervised learning. ReaL's core innovation lies in its ability to dynamically manage GPU allocation and parallelization strategies at a fine-grained, task-level during the training process.

The motivation behind ReaL stems from the widespread adoption of RLHF as a critical post-training technique for aligning LLMs with human instructions and preferences, exemplified by algorithms like Proximal Policy Optimization (PPO). While supervised training of LLMs benefits from well-established parallelization techniques, RLHF's dynamic data generation, multiple model instances (actor, critic, reference, reward), and varied computational tasks (generation, inference, training) present unique optimization challenges. ReaL aims to overcome the inefficiencies of fixed resource allocation, such as GPU idle time and sub-optimal parallelization, by introducing a two-step process: an offline search for an optimal execution plan using an MCMC-based engine, followed by a runtime engine that dynamically reallocates parameters and adjusts parallelization strategies on the fly. This sophisticated approach promises substantial speedups and improved resource utilization, making large-scale RLHF training more accessible and efficient.

The significance of ReaL extends beyond just performance gains; it represents a crucial step towards democratizing and scaling advanced LLM alignment techniques. By providing a flexible and efficient system, ReaL enables researchers and practitioners to experiment with and deploy complex RLHF algorithms more effectively. The system's ability to adapt to various RLHF algorithms, including DPO, GRPO, and ReMax, further solidifies its position as a versatile tool in the evolving landscape of LLM development. Ultimately, ReaL contributes to the broader goal of creating more capable and human-aligned AI systems by tackling the fundamental infrastructure challenges of reinforcement learning at scale.

Background

▶ Watch: Introduction to ReaL system and RLHF challenges (0:00)

The landscape of large language models (LLMs) has been profoundly impacted by reinforcement learning techniques, particularly Reinforcement Learning from Human Feedback (RLHF). This post-training process is essential for aligning pre-trained LLMs with human instructions and preferences, making them more useful and safe. Beyond RLHF, reinforcement learning is also being applied to reasoning I/O (R1/O1), where models are trained to "think" before generating final outputs. This talk focuses primarily on RLHF, with a brief mention of reasoning I/O as a follow-up.

A fundamental distinction exists between traditional supervised learning and reinforcement learning paradigms. In supervised learning, the data is static and fixed within a dataset. The data flow is straightforward: batches are drawn, passed through a single model for forward and backward passes, and weights are updated. This simplicity has allowed for the development of highly optimized parallelization strategies, often referred to as "3D parallelization" (data, model, pipeline parallelism), which are widely used in pre-training LLMs.

Reinforcement learning, however, presents a significantly more complex computational graph and data flow. Data is not static; it is dynamically generated by the language model itself through a process called exploration. This involves interactions between various modules, such as actors, critics, and an environment. The algorithms are also continuously evolving, demanding highly flexible systems.

Specifically, in RLHF, particularly with the widely adopted Proximal Policy Optimization (PPO) algorithm, the complexity escalates due to the involvement of four distinct language models:

1. Actor: Generates responses based on prompts.
2. Critic: Estimates value functions associated with states or actions.
3. Reference Model: Computes log likelihoods for comparison.
4. Reward Model: Approximates human preferences, learned from human feedback.

All four of these models are typically transformer-based LLMs, each requiring parallelization. Furthermore, the RLHF training process is not monolithic but comprises three distinct computational stages, each with different characteristics and optimal resource utilization patterns:

1. Generation Stage: The actor model sequentially outputs tokens based on a prompt. This is an inference task, but with sequential token generation.
2. Inference Stage: Full sequences are passed through models (reference, reward, critic) to compute likelihoods, values, or rewards. This is a batch inference task.
3. Training Stage: Weights of the actor and critic models are updated using optimization algorithms (e.g., gradient descent). This involves forward and backward passes.

The naive approach to scaling RLHF training, by applying a single, fixed parallelization strategy across all four LLMs and all computational stages, suffers from significant drawbacks. While simple to implement and potentially having low GPU idle time, it often leads to high communication overhead and sub-optimal performance because the fixed strategy cannot adapt to the varying needs of different models and tasks. For instance, a parallelization strategy optimal for sequential generation might be highly inefficient for batch training.

Alternatively, dedicating distinct sets of devices with fixed, optimized parallelization strategies to each of the four models could reduce communication overhead. However, this often results in substantial GPU idle time, as specific models might only be active during certain stages, leaving their assigned GPUs underutilized. More subtly, a single model, such as the actor, is involved in different stages (generation and training) that have fundamentally different computational profiles. A fixed parallelization strategy for such a model would inevitably be sub-optimal for at least one of its roles, thereby reducing overall training throughput. This inherent complexity underscores the need for a more dynamic and fine-grained approach to resource management in RLHF training.

Key Findings

▶ Watch: PPO complexity: Four distinct LLMs require parallelization (2:50)

The central insight and key finding of the ReaL system is that efficient RLHF training for large language models necessitates a task-level execution plan featuring fine-grained GPU allocation and parallelization strategies. This departs significantly from conventional static or coarse-grained parallelization methods common in supervised learning, which are ill-suited for the dynamic and multi-component nature of RLHF. ReaL's approach is predicated on the understanding that different models (actor, critic, reference, reward) and different computational stages (generation, inference, training) within the RLHF pipeline have unique resource demands and optimal parallelization profiles.

ReaL implements a sophisticated two-step process to achieve this fine-grained control:

1. Offline Search Process: Before any RLHF job commences, ReaL executes an MCMC-based search engine. This engine systematically explores a vast search space to identify the optimal GPU allocation and parallelization strategy for each task and model across all training stages. This pre-computation phase is crucial for determining the most efficient execution plan.
2. Runtime Execution Engine: During the actual training run, ReaL's runtime engine dynamically implements the best-found execution plan. This involves parameter reallocation, where model parameters are efficiently redistributed across GPUs, and parallelization strategies are switched on the fly to match the current task's requirements.

The empirical results demonstrate the profound impact of this dynamic and adaptive strategy. ReaL achieved substantial performance improvements, delivering 1.4x to over 3x speedup compared to existing open-source reinforcement training systems. These benchmarks were established on large-scale setups, involving 100 H100 GPUs and training a 70 billion parameter actor model, highlighting ReaL's effectiveness in high-performance computing environments.

Furthermore, the research revealed the superiority of the search-based optimization over human intuition. Comparisons between the MCMC-based execution plan and heuristic hand-implemented execution plans showed that ReaL's automated search engine consistently yielded better execution plans. This improvement was particularly pronounced for scenarios involving longer context lengths, where the complexities of optimal resource management become even more critical.

A significant finding regarding ReaL's versatility is its broad compatibility. The system is designed to be agnostic to the specific RLHF algorithm, capable of supporting various methods such as DPO (Direct Preference Optimization), GRPO (Generalized Reinforcement with Policy Optimization), and ReMax. Users can define their RLHF algorithm as a dependency graph, which the search engine then processes to generate an optimized execution plan, ensuring speedups across a spectrum of advanced RLHF techniques. This flexibility underscores ReaL's potential as a universal platform for efficient LLM alignment research and deployment.

Technical Deep Dive

▶ Watch: Three different computational stages in RLHF training (4:00)

The technical foundation of ReaL lies in its sophisticated approach to managing the inherent complexity of RLHF training, which involves multiple large language models and distinct computational phases. The core problem ReaL solves is the sub-optimality arising from fixed GPU allocation and parallelization strategies. Traditional parallelization, while effective for supervised training, struggles with RLHF's dynamic data generation, varied model roles (actor, critic, reference, reward), and diverse computational tasks (generation, inference, training).

ReaL introduces a task-level execution plan that orchestrates dynamic GPU allocation and parallelization strategies. This plan dictates not only which devices each model is placed on but also the specific parallelization technique (e.g., data parallelism, model parallelism, pipeline parallelism) to be used for a given model during a particular computation stage. For instance, in the inference stage, the plan might specify that reference computation and reward computation run concurrently, while the critic computation utilizes all available GPUs. In the training stage, the plan could allocate half the GPUs to train the actor and the other half to train the critic. Crucially, a single model, like the actor, might employ different parallelization strategies during its generation stage versus its training stage, recognizing that optimal strategies for these tasks can be vastly different.

The implementation of such a dynamic plan necessitates parameter reallocation. This involves redistributing the parameters of language models across GPUs and changing the underlying parallelization strategies on the fly during the training process. While this introduces an overhead (represented as "yellow bars" in the speaker's visualization), ReaL is designed to minimize this overhead through smart implementation, making it a viable trade-off for overall efficiency gains. The benefit of this dynamic approach is a significant reduction in GPU idle time and communication overhead, leading to higher end-to-end training throughput.

The overall system design of ReaL comprises three main components:

1. Input: The system takes various settings as input. These include training settings (e.g., batch size, number of iterations, number of models), cluster settings (e.g., number of GPUs available), and data flow settings. Critically, any reinforcement training algorithm can be represented as a dependency graph and fed into the system, enabling algorithm agnosticism.

1. Execution Plan Generator: This is the heart of ReaL's optimization capabilities. It employs an MCMC-based search engine to find the optimal execution plan. The process begins with a lightweight profiler. This profiler runs a small batch size through every layer of each model to estimate the runtime and memory consumption for individual layers and models under different configurations. This profiling data is crucial for the MCMC search engine to evaluate the feasibility and efficiency of potential execution plans, especially considering that memory usage is a critical constraint in RLHF training.

The search space for optimal plans is notoriously large, encompassing choices for device placement, parallelization strategy for each of the four LLMs, and concurrent task execution. Given this complexity, a local search method like Markov Chain Monte Carlo (MCMC) is employed. The MCMC engine runs for a fixed duration (e.g., 5-10 minutes) before the main training job starts, to approximate a good plan within this vast search space. The paper includes detailed analysis of this MCMC algorithm.

1. Runtime Engine: Once the execution plan generator outputs the best-found plan (a computation graph with associated strategies and allocations), it is fed to a centralized master controller. This controller manages all worker nodes, orchestrating the execution of computational jobs according to the plan. This includes dynamically switching parallelization strategies and reallocating parameters as dictated by the current task. A crucial aspect of the runtime engine is the implementation of distributed data storage. Since parameter reallocation and strategy switching can also necessitate shifting prompts or training data across different machines, an efficient distributed data storage mechanism ensures that data movement overhead remains minimal, complementing the parameter reallocation efficiency. During the Q&A, it was confirmed that the system does account for the possibility of overlapping parameter reallocation communication with computation, further minimizing its impact.

In essence, ReaL transforms the static, monolithic view of LLM training into a dynamic, adaptive orchestration of computational resources. By treating parallelization and allocation as mutable, task-dependent variables, it achieves a level of fine-grained control that significantly boosts the efficiency of complex RLHF pipelines.

Experimental Setup & Results

▶ Watch: ReaL's core problem: fine-grained task-level resource management (7:30)

The experimental evaluation of ReaL was conducted using the system version available at the time of the MLSys paper submission, with the disclaimer that systems are continuously evolving. This ensures a fair comparison against contemporary baselines.

Hardware and Model Scale:

The experiments were performed on a substantial cluster comprising 100 NVIDIA H100 GPUs. The primary model under consideration was a 70 billion parameter actor model, representing a significant scale for large language models. This setup allowed for rigorous testing of ReaL's capabilities in a demanding, real-world environment.

Baselines:

ReaL's performance was benchmarked against "all the existing reinforcement training systems" that were open-source at the time of the paper. While specific names of these baseline systems were not detailed in the talk, the implication is a comparison against state-of-the-art open-source solutions for RLHF training.

Headline Results & Speedup:

The most significant finding from the experiments was the substantial speedup achieved by ReaL. The system demonstrated 1.4x to over 3x speedup in end-to-end RLHF training jobs compared to the existing open-source systems. This broad range highlights ReaL's effectiveness across different configurations and potentially varying complexities of RLHF tasks.

Ablation Studies and MCMC Search Efficacy:

A key aspect of the experimental validation involved comparing ReaL's MCMC-based execution plan generator against heuristic hand-implemented execution plans. The results clearly indicated that the search-based engine consistently produced superior execution plans. This improvement was particularly notable when dealing with longer context lengths, suggesting that as the complexity of the task increases, the benefits of automated, optimized search become more pronounced. This validates the design choice of an MCMC search over manual optimization for such a complex problem space.

Compatibility with Diverse RLHF Algorithms:

The experiments also confirmed ReaL's versatility and compatibility with various RLHF algorithms. The system was tested with different algorithms beyond standard PPO, including DPO (Direct Preference Optimization), GRPO (Generalized Reinforcement with Policy Optimization), and ReMax. For all these algorithms, ReaL demonstrated significant speedups, affirming its ability to generalize its optimization approach by converting any reinforcement algorithm into a dependency graph for the search engine. This showcases ReaL as a flexible platform, not just tied to one specific RLHF method.

Further Analysis:

The speaker briefly mentioned that task-level or kernel-level analyses were also performed, with details available in the paper for those interested. This indicates a deeper investigation into the specific bottlenecks and optimizations at granular levels within the system.

In summary, ReaL's experimental results, conducted on a large-scale H100 GPU cluster with a 70B parameter model, convincingly demonstrate its superior efficiency. The 1.4x to 3x speedup, combined with the proven effectiveness of its MCMC-based search over manual heuristics and its broad algorithm compatibility, positions ReaL as a significant advancement in the field of efficient RLHF training for LLMs.

Practical Implications

▶ Watch: ReaL's two-step solution: search for optimal execution plan (8:00)

ReaL's dynamic and fine-grained resource management for RLHF training holds significant practical implications for various stakeholders in the AI/ML ecosystem.

For practitioners and model builders, ReaL means that training large language models with sophisticated RLHF techniques becomes significantly more efficient and accessible. The substantial speedups (1.4x to 3x) translate directly into reduced training times, allowing for faster iteration cycles, more extensive hyperparameter tuning, and the ability to train larger, more capable models within existing computational budgets. This accelerates research and development in LLM alignment, enabling quicker deployment of models that better adhere to human preferences and instructions. Furthermore, ReaL's compatibility with diverse RLHF algorithms (PPO, DPO, GRPO, ReMax) empowers practitioners to experiment with cutting-edge methods without being bottlenecked by system-level inefficiencies, fostering innovation in the field.

For infrastructure teams and deployers, ReaL offers a sophisticated solution to optimize GPU utilization, a critical and expensive resource. By minimizing GPU idle time and reducing communication overhead through dynamic allocation and parallelization, ReaL ensures that computational resources are used as effectively as possible. This leads to cost savings and improved throughput on existing hardware. The system's ability to automatically generate optimal execution plans reduces the manual effort and expertise required from infra teams to configure complex distributed training setups for RLHF, simplifying operations and reducing potential human error. The centralized controller and distributed data storage mechanisms provide a robust framework for managing complex, multi-stage, multi-model training pipelines.

Tradeoffs and Limitations:

While ReaL offers considerable advantages, there are inherent tradeoffs. The primary one is the initial search time for the MCMC-based execution plan. The speaker mentioned this process takes "five or 10 minutes" before the training job starts. For very short training runs or rapid prototyping, this initial overhead might be noticeable, though for large-scale, long-running RLHF jobs, it is a negligible investment for substantial long-term gains.

Another consideration is the overhead of parameter reallocation itself. Although the speaker emphasized that it can be implemented "in a very smart way" to be small, and that the system accounts for overlapping reallocation communication with computation, it is an additional complexity compared to static setups. Ensuring this overhead remains minimal across various hardware configurations and network conditions is crucial for the system's robustness.

The dynamic nature of ReaL, while powerful, also implies a higher degree of system complexity compared to simpler, fixed-strategy approaches. This might require a steeper learning curve for users or a more mature deployment environment. However, the open-source nature of the project and its user-friendly input (converting algorithms to dependency graphs) aim to mitigate this.

In essence, ReaL represents a significant step towards making advanced LLM training, particularly RLHF, more scalable, efficient, and broadly applicable. It tackles the fundamental infrastructure challenges that have historically limited the widespread adoption of complex reinforcement learning techniques in the context of large-scale deep learning, paving the way for more sophisticated and human-aligned AI models.

Key Takeaways

• Complexity of RLHF Training: Unlike supervised learning, RLHF involves multiple interacting language models (actor, critic, reference, reward) and distinct computational stages (generation, inference, training), making efficient parallelization challenging.
• Need for Dynamic Resource Management: Fixed GPU allocation and parallelization strategies lead to significant inefficiencies, including GPU idle time and sub-optimal performance, due to the varied demands of different models and tasks within RLHF.
• ReaL's Two-Step Optimization: The system employs a novel two-step approach: an MCMC-based search engine generates a fine-grained, task-level execution plan offline, which is then dynamically implemented by a runtime engine that reallocates parameters and adjusts parallelization strategies on the fly.
• Significant Performance Gains: ReaL achieves substantial speedups of 1.4x to over 3x compared to existing open-source RLHF training systems, demonstrating superior efficiency on large-scale setups (e.g., 70 billion parameter actor on 100 H100 GPUs).
• Automated Search Outperforms Heuristics: The MCMC-based search engine consistently finds better execution plans than hand-implemented heuristics, especially for complex scenarios like longer context lengths.
• Broad Algorithm Compatibility: ReaL is designed to be algorithm-agnostic, supporting various RLHF methods such as PPO, DPO, GRPO, and ReMax by converting them into dependency graphs for optimization.
• Open-Source Contribution: ReaL is an open-source project, inviting community contributions and broader adoption to advance efficient RLHF training.

About the Speaker(s)

The primary speaker for this talk is Zhiyu Mei, affiliated with Chinghua University and leading a reinforcement learning lab at Ant Research. While the transcript does not provide extensive personal background details, his dual affiliation highlights expertise spanning both academic research and industrial application in the field of reinforcement learning. The presentation reflects a deep understanding of the practical challenges in scaling reinforcement learning to large language models and a commitment to developing open-source solutions.

The talk briefly mentions a follow-up project, A Real, which extends the ReaL system to address reasoning in reinforcement learning. This project, also open-source and developed by the same team, underscores Ant Research's active involvement in pushing the boundaries of RL for LLMs. The speaker also mentioned that Ant Research is actively hiring for positions in Sunnyvale, Beijing, and other offices, indicating a growing team and continued investment in this area. The other listed speakers, Wei Fu, Kaiwei Li, Guangju Wang, Huanchen Zhang, and Yi Wu, are co-authors and likely key contributors to the ReaL project, reflecting a collaborative effort from the research group.