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
分类: cs.LG, cs.AI
发布日期: 2026-03-24
💡 一句话要点
SortedRL通过在线长度感知调度加速LLM的强化学习训练
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 强化学习 大型语言模型 在线调度 长度感知 训练加速 rollout优化 推理能力
📋 核心要点
- 现有RL方法在训练LLM时,rollout阶段耗时过长,成为性能瓶颈,尤其是在生成长序列时,自回归生成和同步开销巨大。
- SortedRL的核心思想是根据rollout样本的长度进行排序和调度,优先处理短样本,从而实现更高效的训练和更好的on-policy特性。
- 实验结果表明,SortedRL能够显著降低训练bubble ratios,并在多个逻辑和数学任务上提升LLM的性能,最高提升达18.4%。
📝 摘要(中文)
为了提升大型语言模型(LLM)的推理能力,尤其是需要长链思考生成的任务,扩展强化学习(RL)展现出强大的潜力。然而,RL训练效率通常受限于rollout阶段,当生成长轨迹(例如16k tokens)时,由于自回归生成速度慢以及rollout和策略更新之间的同步开销,rollout阶段可能占据总训练时间的70%。我们提出了SortedRL,一种在线长度感知调度策略,旨在通过提高rollout效率和保持训练稳定性来解决这一瓶颈。SortedRL根据输出长度重新排序rollout样本,优先处理短样本形成组以进行早期更新。这同时实现了大型rollout批次、灵活的更新批次和近on-policy的微课程构建。为了进一步加速pipeline,SortedRL结合了一种通过基于缓存的机制来控制off-policy训练程度的机制,并由专门的RL基础设施支持,该基础设施通过有状态的控制器和rollout缓冲区管理rollout和更新。在使用LLaMA-3.1-8B和Qwen-2.5-32B在包括逻辑谜题和数学挑战(如AIME 24、Math 500和Minerval)等各种任务上的实验表明,SortedRL将RL训练bubble ratios降低了50%以上,同时在相同数据量下获得了比基线高3.9%到18.4%的性能。
🔬 方法详解
问题定义:现有强化学习方法在训练大型语言模型时,rollout阶段的效率低下是一个主要瓶颈。特别是当需要生成长序列时,自回归生成过程缓慢,并且rollout和策略更新之间的同步开销巨大,导致训练时间显著增加。现有方法难以兼顾rollout效率和训练稳定性,限制了LLM在复杂推理任务上的性能提升。
核心思路:SortedRL的核心思路是根据rollout样本的长度进行在线排序和调度。通过优先处理短样本,可以更快地进行策略更新,从而减少训练过程中的空闲时间(bubble ratios)。同时,通过控制off-policy训练的程度,可以维持训练的稳定性。这种长度感知的调度策略旨在优化rollout效率,并构建一个近似于on-policy的微课程,从而加速LLM的强化学习训练。
技术框架:SortedRL的整体框架包括以下几个主要模块:1) Rollout Buffer:用于存储rollout生成的样本。2) Stateful Controller:负责管理rollout和更新过程,并根据样本长度进行排序和调度。3) Length-Aware Scheduler:根据样本长度动态调整rollout和更新的批次大小。4) Cache-Based Off-Policy Control:通过缓存机制控制off-policy训练的程度,以维持训练稳定性。整个流程通过一个专门的RL基础设施进行管理,确保高效的rollout和更新。
关键创新:SortedRL的关键创新在于其在线长度感知调度策略。与传统的随机或固定顺序的rollout方法不同,SortedRL根据样本长度动态调整rollout和更新的顺序,从而优化训练效率。此外,SortedRL还引入了一种基于缓存的机制来控制off-policy训练的程度,这有助于维持训练的稳定性。这种长度感知调度和off-policy控制的结合是SortedRL与现有方法的本质区别。
关键设计:SortedRL的关键设计包括:1) 长度排序算法:用于根据样本长度对rollout样本进行排序。2) 动态批次大小调整:根据样本长度动态调整rollout和更新的批次大小,以优化训练效率。3) 缓存大小和更新频率:用于控制off-policy训练的程度,需要根据具体任务进行调整。4) 损失函数:采用标准的强化学习损失函数,如PPO或TRPO,并根据需要进行调整。
🖼️ 关键图片
📊 实验亮点
实验结果表明,SortedRL在多个任务上显著优于基线方法。例如,在AIME 24、Math 500和Minerval等数学挑战任务上,SortedRL在相同数据量下获得了3.9%到18.4%的性能提升。此外,SortedRL还能够将RL训练bubble ratios降低50%以上,表明其能够更有效地利用计算资源。
🎯 应用场景
SortedRL具有广泛的应用前景,可用于提升LLM在各种需要长链思考和复杂推理的任务中的性能,例如数学问题求解、逻辑推理、代码生成等。该方法可以显著降低训练成本,加速LLM的开发和部署,并推动LLM在实际应用中的普及。
📄 摘要(原文)
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.