Source-linked AI summary
SortedRL: Accelerating RL Training for LLMs through Online Length-Aware Scheduling
Yiqi Zhang, Huiqiang Jiang, Xufang Luo, Zhihe Yang, Chengruidong Zhang, Yifei Shen, Dongsheng Li, Yuqing Yang, Lili Qiu, Yang You
TL;DR
Long autoregressive rollouts dominate RL training and become inefficient when response lengths vary under synchronized updates. SortedRL uses online length-aware scheduling, controllable off-policy sampling, and stateful rollout infrastructure to address this bottleneck. It reduces rollout bubbles and improves sample efficiency across logical and mathematical reasoning tasks, while interrupted long generations remain more exposed to off-policy effects.
Problem
Variable-length autoregressive rollouts create synchronization bubbles and make generation a major RL training bottleneck for LLMs.
Method
SortedRL sorts rollout samples by output length, enables flexible on-policy or partial off-policy updates, and coordinates them with a length-aware controller and stateful buffer.
Results
74% to less than 5.81%: SortedRL cut rollout bubble ratio while achieving over 18% higher AIME24 accuracy than baseline with identical training data.
Takeaways & Limitations
SortedRL provides a compute- and sample-efficient approach for RL-based LLM tuning across logical reasoning and mathematical tasks.
Takeaways & Limitations
Interrupted longer generations are more affected by off-policy resumption because they undergo more interruptions.
Abstract
from arXiv · showhide
Scaling reinforcement learning (RL) has shown strong promise for enhancing the reasoning abilities of large language models (LLMs), particularly in tasks requiring long chain-of-thought generation. However, RL training efficiency is often bottlenecked by the rollout phase, which can account for up to 70% of total training time when generating long trajectories (e.g., 16k tokens), due to slow autoregressive generation and synchronization overhead between rollout and policy updates. We propose SortedRL, an online length-aware scheduling strategy designed to address this bottleneck by improving rollout efficiency and maintaining training stability. SortedRL reorders rollout samples based on output lengths, prioritizing short samples forming groups for early updates. This enables large rollout batches, flexible update batches, and near on-policy micro-curriculum construction simultaneously. To further accelerate the pipeline, SortedRL incorporates a mechanism to control the degree of off-policy training through a cache-based mechanism, and is supported by a dedicated RL infrastructure that manages rollout and update via a stateful controller and rollout buffer. Experiments using LLaMA-3.1-8B and Qwen-2.5-32B on diverse tasks, including logical puzzles, and math challenges like AIME 24, Math 500, and Minerval, show that SortedRL reduces RL training bubble ratios by over 50%, while attaining 3.9% to 18.4% superior performance over baseline given same amount of data.
1 Introduction
RL improves LLM reasoning, but long autoregressive rollouts and variable response lengths make generation the main training bottleneck. SortedRL addresses this by sorting samples online, controlling off-policy training, and coordinating rollout and updates.
- RL can enhance pretrained LLMs’ reasoning abilities on complex tasks, including mathematical problem solving and competition-level coding.
- Longer generated responses can accompany performance improvements because they often contain richer intermediate reasoning steps.
- Autoregressive long-rollout generation is the primary RL training bottleneck, while variable lengths force updates to wait for the longest response and create hardware bubbles.
- Large rollout batches can reduce idling, but repeatedly updating on fixed-size batches makes data increasingly off-policy and can undermine training stability.
- SortedRL sorts responses by length, updates shorter samples earlier, caches unfinished samples, and coordinates rollout and updates with a length-aware controller and stateful buffer.
- 74% to less than 6%: SortedRL sharply reduced bubble ratio while achieving 3.9% to 18.4% better performance on competition-level mathematics with the same training data.It also reached the same high logical-reasoning score with 40.74% fewer samples.
2 Motivation and Preliminaries
RL for LLMs combines rollout, inference, and model-update stages, but long and heterogeneous generations make synchronized rollout a major source of idle computation. The resulting bubbles limit efficient scaling, especially for long trajectories.
- RL for LLMs: RL training consists of actor rollout, critic/reward/reference-model inference, and gradient-based model updates.
- Policy optimization: PPO and Reinforce++ are REINFORCE-based policy optimization methods whose advantage computation and normalization can differ.
- Policy optimization: The policy objective uses an importance-sampling ratio clipped by ε and an advantage estimator Â_t.
- Rollout bottleneck: With a 16K-token maximum generation length, rollouts can consume 70% of computational resources as actor generation increasingly dominates RL duration.
- Rollout bottleneck: In a sampling batch of 512, 80% of responses finish within 3K tokens while 5% reach the token limit, producing substantial length heterogeneity.
- Rollout bottleneck: Synchronous rollout and actor updates prevent standard continuous batching from eliminating idle periods, and longer trajectories worsen the resulting bubble ratios.
3 SortedRL
SortedRL combines online length-aware scheduling, controllable off-policy sampling, and co-designed infrastructure to reduce rollout bubbles while preserving flexible training behavior. It groups similar-length responses, manages interrupted trajectories, and coordinates updates through stateful control.
- SortedRL components: SortedRL aligns samples with similar generation lengths, supports on-policy and partial off-policy modes, and uses a controller with a stateful rollout buffer.
- Framework: The framework feeds buffered prompts, applies early termination, collects trajectories, updates the model, and sorts training batches.
- Online length-aware scheduling: Oversubscription keeps the rollout engine near its optimal batch size, while threshold-based early termination harvests completed and partial outputs to reduce idle computation.
- Online length-aware scheduling: Interrupted longer generations create a limitation: resumed segments are off-policy, and repeated interruptions affect longer sequences more strongly.
- Grouped rollout: Cache-aware loading consumes grouped prompts before admitting new ones, preventing prompt starvation and maintaining balanced training dynamics.
- Micro-curriculum: Length-sorted responses form an incremental-difficulty micro-curriculum, while selective batching lets the controller feed ready trajectories in an order suited to batch-normalized algorithms.
- Controlled off-policiness: Fully on-policy mode discards unfinished requests, whereas partial mode caches generated tokens and log probabilities so interrupted trajectories can resume with exact prior-token probabilities.
- Infrastructure: The infrastructure tracks unconsumed prompts, scavenged segments, and their log probabilities, while the stateful buffer stores partial trajectories and lifecycle status.
4 Experiment
Experiments evaluate SortedRL on logical and mathematical reasoning tasks, comparing on-policy and partial variants with baseline methods. Results examine learning behavior, benchmark accuracy, rollout throughput, bubble ratios, and sensitivity to grouping size.
- Experiment Setup: The evaluation uses LogicRL logical puzzles and DAPO-Math-17k mathematical problems, with six external mathematics benchmarks for capability assessment.LogicRL contains 5,000 synthetic Knights and Knaves puzzles; DAPO-Math-17k contains integer-answer problems from AoPS, and the benchmarks include GSM8k, MATH500, Minerva Math, OlympiadBench, AIME 2024, and AMC 2023.
- Results on Logic Problems: On LogicRL, on-policy SortedRL improved evaluation faster than baseline and began exploring lengthy responses 150 steps earlier.The partially off-policy variant experienced an explosion in reasoning length followed by unrecoverable performance degradation.
- Results on Math Problems: After 600 update steps on AIME24, baseline, on-policy SortedRL, and partial SortedRL achieved 23.33%, 20.83%, and 19.69% accuracy (mean@32), respectively.Across MATH500, Minerva, Olympiad, and AIME24, token-efficiency and performance followed decreasing off-policiness, while GSM8K showed an inverse pattern.
- Results on Math Problems: SortedRL provides two relatively on-policy methods while allowing the system to operate at the maximum batch.The modes differ in their degree of off-policy training, enabling comparison between fully on-policy and partial operation.
- Throughput of Different Methods: Rollout throughput reached 3987, 4289, and 5559 output tokens per second for baseline, fully on-policy, and partial modes, respectively.Relative to the baseline bubble ratio of 74%, on-policy and partial modes reduced it to 5.81% and 3.37%.
5 Related Work
RLHF frameworks have evolved from algorithm-centric libraries to scalable, throughput-oriented systems. Recent systems integrate optimized LLM-serving techniques to accelerate rollout execution.
- RLHF infrastructure progressed from algorithm-centric libraries to throughput-oriented systems that scale to thousands of GPUs.
- Modern RLHF systems use high-throughput serving components such as PagedAttention, Radix Attention, optimized kernels, graph capture, continuous batching, and chunked prefill.
6 Conclusion
SortedRL is an online length-aware scheduling scheme for RL-based LLM tuning that targets hardware utilization, sample efficiency, and rollout bubbles. Experiments report improved sample efficiency, substantially lower bubble ratios, and reduced off-policy effects.
- SortedRL provides online length-aware scheduling that reorders training batches to construct sample-efficient micro-curricula.
- Over 18% accuracy improvement was achieved on AIME24 with Qwen-2.5-32B using the same training-data amount as the baseline.
- Bubble ratio fell from 74% to less than 5.81%, producing up to nearly 40% higher rollout throughput.
- SortedRL offers fully on-policy and partial off-policy methods with less off-policy effect than canonical large-batch training.
A Dataset Example
The appendix presents examples from the LogicRL logical-puzzle task and a mathematical-reasoning task. The examples show prompts, expected answer formats, and a concrete geometry problem.
- The LogicRL example includes a ground-truth assignment identifying Michael, Zoey, and Ethan as knights.
- LogicRL uses Knights and Knaves puzzles in which players deduce character roles from statements.
- The mathematical-reasoning example instructs the model to solve a problem step by step and place the answer on its own line.
- The example geometry problem asks for a+b+c from a constrained triangle-and-angle construction with AB+AC expressed as a√b/c.
B Additional Analysis
Additional analysis links response length with training behavior and examines local and aggregate performance patterns under SortedRL. The results show recurring short-to-long rollout structure and length-dependent validation dynamics.
- Response length is highly associated with model performance and tends to increase alongside reasoning capability.
- SortedRL rollout batches exhibit a short-short-short-long pattern that forms an online micro-curriculum.
- Longer responses tend to have lower local performance, but the iteration after a long batch can achieve a higher score than preceding short batches.
- Performance improvements often appear on the falling edge of response length, while SortedRL models have longer responses at comparable validation performance.
C Author Contributions
The authors jointly initiated the project, designed and guided SortedRL’s overall framework, and contributed to algorithm design and analysis.
- All authors initiated the project, designed the overall framework, guided SortedRL’s development, and contributed to algorithm design and analysis.