通过动态排序生成长度加速大模型强化学习训练
SortedRL: Accelerating RL Training for LLMs through Online Length-Aware Scheduling
- 按输出长度动态排序样本,优先处理短文本以提升并行效率
- 训练耗时减少50%以上,相同数据下性能提升3.9%至18.4%
- 适合长链推理任务的RL训练,尤其适用于数学与逻辑类挑战
强化学习(RL)在提升大语言模型(LLMs)推理能力方面展现出巨大潜力,尤其在需要长思维链生成的任务中。然而,由于自回归生成慢且存在同步开销,生成长轨迹(如16k token)时,推理阶段可占总训练时间高达70%。本文提出SortedRL,一种在线长度感知调度策略,通过按输出长度重排推理样本,优先处理短样本并分组提前更新,实现大规模推理批处理、灵活更新批处理和近似在线微课程。为进一步加速流水线,引入基于缓存的离策略训练控制机制,并配备专用基础设施,由状态控制器和推理缓冲区管理推理与更新。在LLaMA-3.1-8B和Qwen-2.5-32B上进行实验,涵盖逻辑谜题及数学挑战(如AIME 24、Math 500、Minerval),结果表明SortedRL将训练空转比例降低超50%,在相同数据量下性能优于基线3.9%至18.4%。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。