Source-linked AI summary
Optimizing RLHF Training for Large Language Models with Stage Fusion
Yinmin Zhong, Zili Zhang, Bingyang Wu, Shengyu Liu, Yukun Chen, Changyi Wan, Hanpeng Hu, Lei Xia, Ranchen Ming, Yibo Zhu, Xin Jin
TL;DR
Existing RLHF frameworks miss optimization opportunities within tasks, while generation suffers from long-tailed samples and training from pipeline bubbles. RLHFuse splits work into finer-grained subtasks and fuses stage execution; it increases throughput by up to 3.7× compared to state-of-the-art solutions.
Problem
Existing RLHF frameworks treat workflows as compositions of tasks, missing internal optimization opportunities amid generation-stage long tails and training-stage pipeline bubbles.
Source: §1 · pp. 1–2Method
RLHFuse breaks RLHF tasks into finer-grained subtasks and uses inter- and intra-stage fusion to improve training throughput.
Source: §1 · p. 2Results
Throughput improves by up to 3.7× compared to state-of-the-art solutions.
Source: §1 · p. 2Takeaways & Limitations
Subtask-level fused execution mitigates generation data skewness and training pipeline bubbles while increasing GPU utilization.
Source: §9 · p. 12One migration is usually sufficient when long-tailed sample generation exceeds inference-task execution time.
Source: §4.2 · p. 6
Abstract
from arXiv · showhide
We present RLHFuse, an efficient training system with stage fusion for Reinforcement Learning from Human Feedback (RLHF). Due to the intrinsic nature of RLHF training, i.e., the data skewness in the generation stage and the pipeline bubbles in the training stage, existing RLHF systems suffer from low GPU utilization. RLHFuse breaks the traditional view of RLHF workflow as a composition of individual tasks, splitting each task into finer-grained subtasks, and performing stage fusion to improve GPU utilization. RLHFuse contains two key ideas. First, for generation and inference tasks, RLHFuse splits them into sample-level subtasks, enabling efficient inter-stage fusion to overlap the execution of generation and inference stages, thus mitigating the original generation bottleneck dominated by long-tailed samples. Second, for training tasks, RLHFuse breaks them into subtasks of micro-batches and performs intra-stage fusion to concurrently execute these subtasks in the training stage with a fused pipeline schedule, effectively mitigating the pipeline bubbles. The experiments show that RLHFuse increases the training throughput by up to $3.7\times$, compared to existing systems.
1 Introduction
Existing RLHF frameworks miss optimization opportunities by treating workflows as individual tasks, despite long-tailed generation and training pipeline bubbles. RLHFuse instead fuses finer-grained subtasks across stages.
1 Introduction
- Long-tailed generation leaves a few samples running while later inference waits, reducing GPU utilization.The issue grows more significant as maximum output lengths increase.Source: §1 · pp. 1–2
- Training pipeline bubbles reduce efficiency as models and pipeline parallelism scale, with remaining bubbles challenging synchronous training.Source: §1 · p. 2
- RLHFuse splits generation and inference into sample-level subtasks, migrating remaining long-tail samples so freed GPUs can run inference.For training, it uses micro-batch subtasks and a fused schedule to fill pipeline bubbles.Source: §1 · p. 2
- Up to 3.7× higher throughput than state-of-the-art solutions is reported for RLHFuse.The system also applies optimizations across RLHF stages to support internal-product training.Source: §1 · p. 2
2 Motivation
The section introduces the RLHF workflow and explains why its generation and training stages underuse GPUs. Long-tailed outputs delay inference, while pipeline bubbles leave training resources idle.
2.1 Background
- The focused PPO alignment stage generates Actor responses, runs inference on them, then trains Actor and Critic on mini-batches.After Actor training, its updated parameters generate samples for the next iteration.Source: §2.1 · pp. 2–3
- An RLHF iteration uses Actor, Reference, Reward, and Critic models; the Reference and Reward models are frozen, while the Critic supplies action-level feedback.Source: §2.1 · pp. 2–3
2.2 Problems in RLHF Training
- Long-tail responses block inference until generation finishes, so even a few unusually long samples can lower GPU utilization.Decoding is memory-bandwidth-bound and typically needs batches of a few hundred for high utilization.Source: §2.2 · p. 3
- In LMSYS-Chat-1M, the P99.9 output length exceeds ten times the median, and long-tail generation accounts for more than half of generation time in the internal-model breakdown.The figure shows output-length CDFs and iteration-time breakdowns under different maximum output lengths.Source: §2.2 · p. 3, Figure 2 · p. 3
- When pipeline stages scale toward the micro-batch count, 1F1B bubbles can reach about 50%, leaving about half the GPUs idle.Interleaved 1F1B reduces bubbles but adds K-fold communication overhead and can worsen stage imbalance.Source: §2.2 · p. 4, Figure 3 · p. 4
- Existing systems treat tasks as the smallest execution units, overlooking their internal structure and leaving optimization opportunities unused.Source: §2.2 · pp. 3–4
3 RLHFuse Overview
RLHFuse configures parallel and deployment strategies for each task, then fuses execution across generation and inference and within training. Its two fusion mechanisms target generation tails and training pipeline bubbles.
3 RLHFuse Overview
- RLHFuse first assigns task-specific parallel strategies, then launches tasks according to workflow dependencies while handling weight redistribution and data transfer.Source: §3 · p. 4
- During generation, RLHFuse migrates remaining samples to dedicated instances and repurposes the original instances to run inference concurrently.This overlaps inference with generation of long-tail samples.Source: §3 · pp. 4–5
- During training, a lightweight algorithm builds a fused Actor–Critic pipeline schedule based on model sizes and parallel strategies, reducing bubbles and activation memory.Source: §3 · p. 5
4 Data-aware Inter-Stage Fusion
RLHFuse overlaps generation and inference by refining their dependency from the stage level to the sample level, then migrates remaining long-tail samples to free GPUs for inference. Its fused execution plan balances migration timing, destination capacity, and migration cost.
4.1 Opportunities and Challenges
- Sample-level dependencies let inference begin as soon as generation finishes for individual samples, without changing synchronous training semantics.The fused plan detects the long-tail decoding phase, migrates remaining samples to a few instances, and uses freed resources to start inference early.Source: §4.1 · p. 5
- Inter-stage fusion must choose migration timing, destination, and mechanism carefully because each affects the two stages’ overall execution time.The objective is to maximize overlap without changing the original generation task’s execution time.Source: §4.1 · p. 5, §4.2 · p. 5
4.2 Fused Execution Plan
- Migration timing trades off generation slowdown from triggering too early against lost overlap from triggering too late.The threshold Rt triggers migration when remaining generation samples fall below a set value; offline simulation tests thresholds from 5% to 95% of the global batch size, and runtime samples can refine the estimate.Source: §4.2 · pp. 5–6
- One migration is usually sufficient when long-tail generation takes longer than inference, allowing the inference stage to be fully overlapped.Repeated migration could enable finer-grained overlap, but is usually unnecessary under this observed timing condition.Source: §4.2 · p. 6
- The destination count must satisfy both a generation-throughput constraint and a key-value-cache memory constraint.The final number of generation instances is set by the stricter of the two constraints; the selected instances are those with the most samples remaining, minimizing migrations.Source: §4.2 · p. 6
- Migration can transfer the key-value cache or send only generated tokens and recompute the prefill; the choice depends on GPU hardware and network bandwidth.With high-bandwidth RDMA in the experiments, cache transfer was chosen and its overhead was negligible compared with iteration time.Source: §4.2 · p. 6
5 Model-aware Intra-Stage Fusion
RLHFuse fuses Actor and Critic micro-batch training tasks in a bidirectional pipeline to fill one another’s bubbles without model replication. Its schedule-generation method handles heterogeneous model sizes and parallel configurations while optimizing latency and activation memory.
5.1 Opportunities and Challenges
- Actor and Critic micro-batches can fill one another’s pipeline bubbles through bidirectional execution, without extra model replication.Chimera’s approach replicates one model and can balance activation memory; RLHF has two distinct models available for fusion.Source: §5.1 · p. 6
- Unlike Chimera’s symmetric schedule, RLHF’s Actor and Critic differ in size and optimal parallel configuration, requiring a schedule for heterogeneous models.Figure 6 contrasts Chimera’s four-stage, four-micro-batch example with fusion of models having different pipeline-stage counts.Source: §5.1 · p. 6, Figure 6 · p. 7
5.2 Fused Pipeline Schedule
- When tensor-parallel degrees differ by a factor s, RLHFuse groups s consecutive stages of one model so both models’ stages use equal GPU resources.The grouping requires the relevant pipeline-stage count to be divisible by s; the transformed problem fuses coprime numbers of pipeline groups.Source: §5.2 · p. 7
- A valid fused schedule preserves micro-batch dependencies, avoids deadlocks, and stays within each stage’s activation-memory limit while minimizing latency.The formulation represents scheduled subtasks as a matrix and prefers lower peak activation memory among schedules with equal latency.Source: §5.2 · p. 7
- The greedy schedule can miss globally useful delays and does not optimize memory, so RLHFuse uses simulated annealing to search for lower-latency schedules.The greedy solution initializes the search; simulated annealing probabilistically explores neighboring schedules and can scale across CPU cores and extend to more models.Source: §5.2 · p. 8
- Schedule latency is computed from each subtask’s inter-stage and intra-stage dependencies, with the maximum finish time across stages defining execution time.The memoized recursion has time complexity O(N × (M1 + M2)); a second search can reduce peak activation memory without degrading latency.Source: §5.2 · p. 8
6 Implementation
RLHFuse extends Megatron-LM with asynchronous device meshes, simulators, and execution support for its two fusion strategies. It also incorporates system optimizations for generation, inference, and training.
6 Implementation
- RLHFuse extends Megatron-LM with multiple device meshes to launch tasks asynchronously using tailored parallelism and deployment strategies.The implementation comprises 7K lines of Python, C++, and CUDA.Source: §6 · p. 8
- The system simulates runtime and memory costs, then searches pruned parallel-strategy configurations for each RLHF task.It follows a model-then-optimize approach similar to ReaLHF and models generation and inference as simplified forward-pass workloads.Source: §6 · pp. 8–9
- A workflow-level simulator models dependencies and communication overhead to select the migration threshold minimizing fused generation-and-inference execution time.The simulation takes the chosen parallel strategies, workload pattern, and threshold Rt as inputs.Source: §6 · p. 9
- MPI parallelizes independent simulated-annealing searches across random seeds, after which the system generates per-device NCCL operations for runtime execution.The independent searches can scale to hundreds of CPU cores.Source: §6 · p. 9
- System optimizations include optimized autoregressive generation, matrix-multiplication-based Generalized Advantage Estimation, and sequence-length balancing across data-parallel groups.The system also reduces Actor and Critic weight-redistribution communication and keeps unchanged Reference and Reward models in CPU memory until needed.Source: §6 · p. 9
- RLHFuse-Base provides an evaluation baseline with the system optimizations but without inter-stage or intra-stage fusion.Source: §6 · p. 9
7 Evaluation
The evaluation measures end-to-end RLHF training throughput and analyzes how stage fusion affects iteration time, latency, and activation memory. Across tested settings, RLHFuse improves throughput over existing systems, while its fusion techniques address generation and training bottlenecks.
7 Evaluation
- Evaluation covers LLaMA models from 13B to 65B on HH-RLHF, four Actor/Critic size pairs, and a 32-node, 256-GPU cluster.Sample throughput is measured in samples per second over 20 consecutive post-warm-up iterations.Source: §7 · p. 9
7.1 End-to-End Results
- Across generation settings with maximum lengths of 512, 1024, and 2048, RLHFuse achieves 2.5×–3.7× higher throughput than DSChat, 1.4×–2.4× than ReaLHF, and 1.2×–1.4× relative improvement over RLHFuse-Base.Source: §7.1 · p. 10
7.2 Performance Analysis
- Inter-stage fusion yields 1.2×–1.6× speedup for generation plus inference, while intra-stage fusion reduces training execution time by 1.2×–1.3×.Other overhead is below 3% of iteration time, and inter-stage migration overhead is negligible.Source: §7.2 · p. 10, Figure 8 · p. 11
7.3 Effectiveness of Stage Fusion
- About 20% of the batch remaining gives the lowest measured fused generation-plus-inference latency, but the migration threshold requires periodic adjustment as output lengths change.The threshold experiment was conducted during early training, and adjustments are usually infrequent because distribution changes are not too drastic.Source: Figure 9 · p. 11, §7.3 · p. 11
- Across tested settings, RLHFuse outperforms baseline methods and reaches the estimated theoretical latency lower bound except in the last case.The lower bound estimates the earliest possible completion time across pipeline stages; a schedule attaining it need not exist.Source: Table 3 · p. 11, §7.3 · p. 11
- For fused schedules, serial 1F1B is the activation-memory lower bound; RLHFuse often reaches it and otherwise keeps extra overhead acceptable.The method also outperforms the greedy approach on memory cost.Source: Table 3 · p. 11, §7.3 · p. 11
7.4 Deep Dive
- In 65B/33B, one 65B model (16 PP stages) fuses with two 33B models (8 each) in reversed directions, fully overlapping 33B training at 65B 1F1B latency.Peak activation memory also reaches the lower bound of executing both models serially with 1F1B.Source: Figure 10 · p. 12, §7.4 · p. 12
8 Related Work
Prior work addresses RLHF model placement, task-specific parallelism, generation speed, and training throughput. The cited approaches do not consider subtask-level opportunities within RLHF tasks, or are orthogonal to RLHFuse's focus.
RLHF training systems
- RLHF frameworks use distributed model placement, hybrid programming, task-specific parallel strategies, or lightweight context switching, but the cited systems do not address subtask-level optimization.Examples include OpenRLHF [38], HybridFlow [5], ReaLHF [6], and PUZZLE.Source: §8 · p. 12
LLM generation optimizations
- Generation systems target throughput, bursty workloads, head-of-line blocking, or prefill-decoding interference; the cited approaches are described as orthogonal to RLHFuse.The passage discusses AlpaServe, FastServe [29], Splitwise, and DistServe [28].Source: §8 · p. 12
LLM training optimizations
- General LLM training work studies pipeline parallelism, automatic parallel-strategy generation, large-scale systems, and multimodal heterogeneity, but is described as agnostic to RLHF characteristics.The cited examples include Alpa, MegaScale, and DistTrain [52].Source: §8 · p. 12
9 Conclusion
RLHFuse uses finer-grained subtasks to fuse execution across and within stages. The evaluation reports increased GPU utilization and training throughput gains of up to 3.7× over existing systems.
9 Conclusion
- RLHFuse views the RLHF workflow at subtask level to enable inter-stage and intra-stage fused execution.The stated targets are data skewness and pipeline bubbles in existing systems.Source: §9 · p. 12
- The evaluation reports that RLHFuse significantly increases GPU utilization.The conclusion does not specify a separate utilization value.Source: §9 · p. 12
- Training throughput increases by up to 3.7× compared with existing systems.Source: §9 · p. 12