用自适应猜测解码加速推理型强化学习训练,不损失精度。
Taming the Long-Tail: Efficient Reasoning RL Training with Adaptive Drafter
- 通过动态调整的轻量级猜测模型,在空闲显卡上持续训练以对齐目标模型。
- 在真实场景下实现1.7倍以上训练速度提升,且保持模型准确率不变。
- 适合需要高效训练大模型推理能力的研究者与工业用户。
大型语言模型(LLM)具备强大推理能力,推动复杂问题求解新进展。然而,使用强化学习(RL)训练这类模型时面临严重效率瓶颈:生成响应呈现显著长尾分布,少数极长响应占据大量计算时间,造成资源浪费与成本上升。为此,我们提出TLT系统,通过集成自适应推测解码实现无损加速。将推测解码应用于RL极具挑战,因工作负载动态变化、目标模型不断演进及猜测模型训练开销大。TLT通过两个协同组件克服难题:(1) 自适应猜测器(Adaptive Drafter),一个在长尾生成期间利用空闲GPU持续训练的轻量级猜测模型,无需额外成本即可保持与目标模型对齐;(2) 自适应回放引擎(Adaptive Rollout Engine),维护内存高效的预捕获CUDAGraph池,并为每批输入自适应选择最优推测策略。评估表明,TLT相比现有最优系统实现超过1.7倍的端到端训练速度提升,保持模型精度,并产出高质量猜测模型,可直接用于高效部署。代码已公开于https://github.com/mit-han-lab/fastrl。
原文摘要 · Abstract (English)
The emergence of Large Language Models (LLMs) with strong reasoning capabilities marks a significant milestone, unlocking new frontiers in complex problem-solving. However, training these reasoning models, typically using Reinforcement Learning (RL), encounters critical efficiency bottlenecks: response generation during RL training exhibits a persistent long-tail distribution, where a few very long responses dominate execution time, wasting resources and inflating costs. To address this, we propose TLT, a system that accelerates reasoning RL training losslessly by integrating adaptive speculative decoding. Applying speculative decoding in RL is challenging due to the dynamic workloads, evolving target model, and draft model training overhead. TLT overcomes these obstacles with two synergistic components: (1) Adaptive Drafter, a lightweight draft model trained continuously on idle GPUs during long-tail generation to maintain alignment with the target model at no extra cost; and (2) Adaptive Rollout Engine, which maintains a memory-efficient pool of pre-captured CUDAGraphs and adaptively select suitable SD strategies for each input batch. Evaluations demonstrate that TLT achieves over 1.7x end-to-end RL training speedup over state-of-the-art systems, preserves the model accuracy, and yields a high-quality draft model as a free byproduct suitable for efficient deployment. Code is released at https://github.com/mit-han-lab/fastrl.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。