FlashRL:开源强化学习加速利器,极速Rollout兼顾性能突破训练瓶颈

FlashRL开源,创新8位量化加速RL大模型训练,结合TIS技术,显著提升效率且不失性能。

原文标题:让强化学习快如闪电:FlashRL一条命令实现极速Rollout,已全部开源

原文作者:机器之心

冷月清谈:

在大型语言模型(LLM)的强化学习(RL)训练中,卷展(rollout)生成是主要的性能瓶颈,占据了总训练时间的大部分。清华AIR和字节跳动SIA Lab团队为此推出了开源解决方案FlashRL,旨在大幅提升RL训练效率。FlashRL的核心在于将8位量化技术(INT8/FP8)应用于rollout阶段,这通常会导致性能下降。然而,该团队通过引入截断重要性采样(TIS)技术,成功解决了量化rollout与高精度训练之间的不匹配问题,使得8位量化在保持甚至超越BF16精度的性能下,实现了显著的吞吐量提升。FlashRL对于模型规模超过140亿参数的大模型效果尤为明显,最高可达1.75倍的吞吐量加速。在内存受限环境下,加速效果更为惊人,某些配置下甚至能实现5倍以上的速度提升。此外,FlashRL还提供了Flash-LLM-RL包,兼容vLLM并支持在线量化。其安装和使用都极其简便,只需一条命令即可集成到现有RL训练流程中,且已证实能在大模型如DAPO-32B上实现端到端加速,同时不牺牲下游任务性能。

怜星夜思:

1、文章里提到,解决强化学习训练中的“rollout生成”瓶颈是FlashRL的关键。为什么rollout会成为如此大的瓶颈呢?除了本文说的计算量大,还有哪些因素让它拖慢整体训练进度?以及,Rollout速度的提升,除了让RL训练更快,还会对强化学习的应用场景带来哪些想象空间?
2、FlashRL的TIS(截断重要性采样)技术解决了量化rollout与高精度训练之间的“不匹配”问题。在更广泛的AI领域,比如推理部署、模型微调或者迁移学习中,是否也存在类似的数据分布或模型精度不匹配问题?如果有,业界通常会用哪些方法来处理这些问题呢?
3、FlashRL强调在大模型上效果更显著(推荐140亿参数以上),那对于中小模型或者资源有限的个人开发者来说,8位量化是否有类似的性价比收益?或者是否存在其他更适合中小模型的优化策略,既能加速又能兼顾性能和资源?

原文内容

机器之心报道

编辑:冷猫


在今年三月份,清华 AIR 和字节联合 SIA Lab 发布了 DAPO,即 Decoupled Clip and Dynamic sAmpling Policy Optimization(解耦剪辑和动态采样策略优化)。这是一个可实现大规模 LLM 强化学习的开源 SOTA 系统,使用该算法,该团队成功让 Qwen2.5-32B 模型在 AIME 2024 基准上获得了 50 分,



  • 论文地址:https://dapo-sia.github.io/static/pdf/dapo_paper.pdf

  • 代码地址:https://github.com/volcengine/verl/tree/gm-tyx/puffin/main/recipe/dapo


中国科学技术大学校友,伊利诺伊大学香槟分校博士,微软研究院的首席研究员、清华大学校友,加州大学圣地亚哥分校计算机科学与工程学院博士生姚峰团队在强化学习的研究中更进一步。


该团队发现,在 DAPO-32B 中,rollout 生成是强化学习训练的主要瓶颈,占据了约 70% 的总训练时间。因此,该团队从 rollout 阶段着手,将 8 bit 量化技术应用于 rollout 生成,并通过 TIS 技术在保持下游性能的同时实现了显著加速。


众所周知,FP8 能让强化学习运行得更快,但往往以性能下降为代价。


刘力源、姚峰团队推出 FlashRL,是首个开源且可用的强化学习实现方案,在推理执行(rollout)阶段应用 INT8/FP8,并且在性能上与 BF16 持平,没有性能损失。该团队在博客中完整发布了该方法的技术细节。



  • 博客标题:FlashRL: 8Bit Rollouts, Full Power RL

  • 博客地址:https://fengyao.notion.site/flash-rl

  • 代码地址:https://github.com/yaof20/Flash-RL


Rollout 量化可能会降低性能


如图 1 和图 2 中 「⋅⋅⋅⋅⋅」 曲线所示,在未使用 TIS 技术的情况下,采用 FP8 或 INT8 进行 rollout 量化,相比 BF16 rollout 会带来显著的性能下降。


这一现象是预期中的,因为 rollout–训练之间的差异被放大了:rollout 是从量化策略 π_int8 采样的,但梯度却是基于高精度策略 π_bf16 计算的。



这种不匹配会使强化学习过程更加偏离策略,从而削弱强化学习训练的有效性。


图 1  左图:吞吐量加速比。FP8 结果在 H100 上测试;INT8 结果分别在 H100 和 A100 上测试。结果基于不同的响应长度和设备测得。右图:Qwen2.5-32B 模型在使用 BF16 rollout 与 INT8 rollout 时的 AIME 准确率对比。所有实验均采用 BF16 FSDP 训练后端。


FlashRL 的独门秘诀


FlashRL 是首个开源且可用的强化学习方案,能够在不牺牲下游性能的前提下使用量化 rollout。


那么,它的「独门秘诀」是什么呢?


解决 Rollout–训练不匹配问题


该团队引入了截断重要性采样(Truncated Importance Sampling,TIS)来减轻 rollout 与训练之间的差距。正如图 1 和图 2 中的实线所示,TIS 使量化 - rollout 训练的性能达到了与采用 TIS 的 BF16 rollout 训练相同的水平 —— 甚至超过了未使用 TIS 的朴素 BF16 rollout 训练。


作者团队之前发表过有关 TIS 的技术博客,感兴趣的读者可以参考:



  • 博客标题:Your Efficient RL Framework Secretly Brings You Off-Policy RL Training

  • 博客链接:https://fengyao.notion.site/off-policy-rl


在这里简单展示一下 TIS 的工作原理。



支持在线量化


现有的推理引擎(如 vLLM)针对大语言模型推理服务进行了优化,但在支持带参数更新的模型量化方面能力有限。该团队提供了 Flash-LLM-RL 包,对 vLLM 进行了补丁,使其能够支持这一功能。


如图所示,FlashRL 的 INT8 可带来高达 1.7 倍的吞吐量提升,同时保持 RL 的优势。此外,如果不使用 TIS 而使用 naive FP8/INT8 ,性能将显著下降。


图 2  左图与中图:在使用量化 rollout 生成的强化学习大语言模型训练中,GSM8K 的准确率表现。请注意,TIS 对缓解分布差异至关重要。右图:π_fsdp 与 π_vllm 之间的 KL 散度。需要注意的是,INT8 rollout 的 KL 散度大于 FP8 rollout 的 KL 散度。


FlashRL 能有多快?


比较在强化学习训练中采用不同 rollout 精度的吞吐量并不简单,因为模型会不断更新,对于同一个查询,不同的量化策略在经过一定的 RL 训练迭代后可能会生成长度不同的回复。


这里将探讨 FlashRL 所实现的加速效果及其对训练效果的影响。


Rollout 加速表现


常规环境下的加速:


研究团队记录了在 7B、14B 和 32B Deepseek-R1-Distill-Qwen 模型上使用 INT8、FP8 和 BF16 精度的 rollout 吞吐量。


图 1 显示了 8 位量化模型相对于 BF16 的加速比。对于较小的 7B 模型,加速比不足 1.2×;而在 32B 模型上,加速比可达 1.75×。这表明量化对大模型的收益远高于小模型。基于分析结果,团队建议仅在模型规模超过 140 亿参数时使用量化


内存受限环境下的加速:


研究团队还评估了在标准推理场景(不涉及 RL)下,采用 8 位量化所能带来的吞吐量提升。具体而言,团队测量了 INT8 的加速比,作为压力测试,用于验证其在 A100/A6000 和 H100 GPU 上的适用性。


使用 vLLM 在相同数据集上分别服务 BF16 与 INT8 量化版本的 Deepseek-R1-Distill-Qwen-32B 模型,并在 A100/A6000 和 H100 GPU 上记录其吞吐量。


图 3  在 4 种仅推理配置下,INT8 量化的 Deepseek-R1-Distill-Qwen-32B 相对于 BF16 的吞吐量加速比,测量结果涵盖不同回复长度。


如图 3 所示,当 GPU 内存成为瓶颈时,量化能够带来极高的加速比 —— 在 TP2-A6000 配置下生成速度提升超过 3 倍,在 TP1-A100 配置下提升甚至超过 5 倍。这突显了量化在 GPU 内存受限场景(如服务更大规模模型)中的巨大潜力。


端到端加速与效果验证


研究团队将 FlashRL 部署于 DAPO-32B 的训练中,以验证所提方法的有效性。由于在图 2 中 FP8 相比 INT8 拥有更小的分布差距,特意选择 INT8 作为更具挑战性的测试场景。


图 4 展示了在 BF16 与 INT8 rollout 下的下游性能与训练加速效果。两种配置在 AIME 基准上的准确率相当,但 INT8 显著提高了训练速度。


这些结果证明,FlashRL 能在不牺牲训练效果的前提下,实现显著的训练加速



图 4. 左图:使用 BF16 与 INT8 rollout 精度进行强化学习训练的下游性能对比。右图:BF16 与 INT8 rollout 在单位小时内可完成的更新步数。所有实验均基于 DAPO 配方,在 Qwen2.5-32B 模型上进行,训练 250 步,硬件配置为 4 个节点、每节点配备 8 张 H100 GPU。


快速使用


使用 FlashRL 只需一条命令! 使用 pip install flash-llm-rl 进行安装,并将其应用于你自己的 RL 训练,无需修改你的代码。


FlashRL 方法支持 INT8 和 FP8 量化,兼容最新的 H100 GPU 以及较老的 A100 GPU。



更多方法细节,请参阅原博客。



© THE END 

转载请联系本公众号获得授权

投稿或寻求报道:liyazhou@jiqizhixin.com

Rollout慢啊,可不就是因为模型得一步一步跟环境互动嘛!你想想,大模型跑一次推理就够吃力了,它还得模拟“我下一步该咋走,走完会咋样”,这个过程循环好多次,数据量还特大,能不慢吗?除了计算,我觉得数据存储和访问也是个问题。生成的那么多经验数据得存起来复用,IO的压力肯定也不小。

要是rollout能飞快,那可太爽了!以前训练个AI下棋得跑好几天,现在可能几个小时就搞定了。这对于游戏AI、虚拟数字人、甚至是一些复杂的金融策略模拟都太重要了。可以更快地尝试不同的策略,让AI变得更聪明,也让我们这些开发者少熬夜。说不定以后AI能自己看几遍攻略就学会玩我们所有游戏了!

那可太有了!AI跑得好不好,除了模型本身,数据适配性、部署环境适配性也是大头。TIS解决的这个问题,其实就是个“训练和部署之间有代沟”的例子。

我们最常碰到的就是“实验室跑得欢,上了生产线就拉胯”的情况。比如,你的AI在干净的数据集上训练得贼好,一到现实世界里那些带噪音、不规范的数据,立马就“懵圈”了。这叫数据分布不匹配

还有就是模型精度不匹配,比如模型用FP32训练的,结果部署的时候为了省资源用了INT8,那性能肯定会掉。解决这事儿的办法可多了:

* 数据增强: 就是把现有数据“变变花样”,模拟更多真实情况,让模型“见多识广”。
* 迁移学习+微调: 大模型先在通用数据集上学个大概,再用我们自己那点儿“小众”数据微调一下,就好比通用技能学会了,再学点专业技能。
* 校准(Calibration): 特别是量化之后,模型可能对某些输出不够“自信”或“过度自信”,得给它校正一下,让它的判断更靠谱。
* 领域适应(Domain Adaptation): 字面意思,让模型能适应新的数据领域,比如你的AI只在白天拍的图片上识别得好,那就得教它怎么识别晚上拍的图片。

说白了,就是得让训练环境和实际应用环境尽量“处对象”,消除彼此的“隔阂”。

关于“rollout为什么是瓶颈”这个问题,从技术原理上讲,rollout本质是一个与环境交互并生成数据(样本)的过程。这个过程是串行的,每次都需要模型进行推理,然后根据推理结果更新环境状态,再进行下一次推理。这不仅涉及大量的模型前向计算,还涉及到CPU与GPU之间的频繁数据传输(例如:环境状态、动作、奖励等),甚至可能是与外部仿真的交互,这些IO操作本身就是耗时的。此外,RL训练通常需要多样化的探索,这意味着需要生成大量rollout数据来覆盖不同的状态-动作空间,这无疑进一步加剧了数据生成的速度需求。

至于rollout速度提升带来的想象空间,我认为这是革命性的。首先,它将极大地缩短RL实验的迭代周期,让研究人员能更快地验证新算法、新模型架构。其次,对于像机器人控制、自动驾驶、复杂系统优化等需要大量真实或高精度仿真环境交互的应用,rollout加速意味着能以更低的成本(时间、计算资源)获取更多高质量的经验,从而训练出更智能、更鲁棒的AI。甚至可能催生出以前因计算成本过高而无法实现的复杂RL应用,比如在超大规模虚拟世界中的AI角色学习,或者极其精细的工业仿真控制。

对于中小模型而言,8位量化(INT8/FP8)仍然能够提供显著的内存占用和功耗优化,这尤其对于资源受限的部署场景(如边缘设备、移动端)具有重要意义。虽然在吞吐量上的加速比可能不像大模型那样惊人(因为小模型本身的计算瓶颈可能不是纯粹的算力,而是缓存、IO等),但综合来看,其能耗和内存优势在多数情况下仍具有高性价比。

针对中小模型的优化策略,除了低比特量化,还有几种常用的方法:
1. 模型结构优化: 设计更高效的神经网络结构,如MobileNet、EfficientNet、ShuffleNet等,这些网络以较少参数和计算量达到可接受的性能。
2. 模型剪枝(Pruning): 移除模型中不重要或冗余的连接/神经元/通道,从而减小模型大小和计算量。剪枝通常需要在剪枝后进行微调以恢复性能。
3. 知识蒸馏(Knowledge Distillation): 使用一个大型、高精度的教师模型来训练一个小型、高效的学生模型。学生模型不仅学习数据的硬标签,还模仿教师模型的输出分布,从而在更小的模型中继承大模型的知识和性能。
4. 混合精度训练(Mixed Precision Training): 在训练阶段就混合使用FP16/BF16和FP32,以加速训练并减少内存占用,但模型最终仍可能部署为FP32或INT8。这与rollout量化是不同层面的优化。

这些策略往往可以组合使用,以达到最佳的性能-资源平衡。

小模型当然也能用量化啊!虽然加速效果可能没大模型那么夸张,但蚊子腿也是肉嘛!对于个人开发者或者资源不那么充裕的团队来说,8位量化最直接的好处就是能让你在显存更小的GPU上跑起来更大的模型,或者让你在同样的GPU上跑更多任务,甚至能省点电费呢!这性价比可一点不低。

当然,除了量化,我们玩小模型的还有很多别的“骚操作”:
* 模型剪枝: 就是把模型里那些“胖子”部分给“减肥”了,减掉不重要的神经元或者连接,模型就轻盈多了。
* 模型蒸馏: 相当于让一个学霸(大模型)把知识“灌输”给一个学渣(小模型),让学渣也能有学霸的水平,但自己学得快,体积也小。
* 用轻量级网络: 比如MobileNet、EfficientNet这些,它们天生就是为移动端或资源有限的环境设计的。

所以说,小模型也有小模型的玩法,一样能玩出花儿来!

哈哈,AI跟人一样,也是个“傲娇”的家伙。它在训练的时候娇生惯养惯了,一旦到了真实世界或者换个“住的地方”(比如从高性能显卡换到低功耗设备),它就可能“水土不服”了,这就是所谓的“不匹配”嘛!

最常见的就两种:
1. 数据水土不服: 比如你训练了个猫狗识别器,只用了室内拍的萌宠照片,结果你把它放到野外去识别,那可能就“抓瞎”了,因为它没见过野外复杂的背景和光线。这时候就得给它做“野外生存训练”(数据增强、领域适应)。
2. 精度“掉链子”: 训练的时候用的是高精度的“精装本”(比如FP32),部署的时候为了省资源,弄了个“盗版漫画”(比如INT8),那画质肯定不行啊,细节都丢了。这时候就得用“鉴宝师”傅永浩(哈哈,开玩笑,是量化感知训练或蒸馏)去修复这些画质,让盗版也能看。

总之,就是要让AI在“训练场”和“战场”的表现尽可能一致,不要在纸上谈兵!

为什么rollout慢?简单说就是“太能生”了,要生大量数据,而且生数据的过程还得“思考人生”(模型推理)。它不仅要计算,还得等环境反馈,这个反馈可能依赖外部系统,所以经常是IO密集型的。想象一下,你生了个娃,还得等他慢慢长大,再培养他,这过程能快吗?

那加速rollout有啥用?最直观的,就是让AI的“迭代速度”加倍,AI可以更快地从错误中学习,从成功中总结经验。这意味着AI能更快地适应新环境,比如在元宇宙里,今天这个规则,明天那个规则,AI能迅速更新它的行为模式。或者在现实世界里,让机器人更快地掌握新技能,比如上午学做咖啡,下午就能冲出世界冠军级的拉花,科幻电影里的场景说不定能提前实现!毕竟,速度就是力量!

是的,在更广泛的AI领域,数据分布或模型精度不匹配是一个非常普遍且关键的问题,尤其是在实际部署中。FlashRL中的TIS解决的是策略(policy)在量化和高精度之间采样分布不一致的问题,这与量化感知训练 (Quantization-Aware Training, QAT) 的目标有些类似,即尽量在训练阶段就模拟量化后的行为,以减少部署时的性能下降。但在更宏观的层面,有几种常见的处理方法:

1. 领域适应(Domain Adaptation):当训练数据与部署数据的领域分布不一致时(A领域训练,B领域使用),常用此法。方法包括特征对齐(如对抗性训练,GAN-based DA)、DANN(Domain Adversarial Neural Networks)、或通过自监督学习在目标领域进行预训练。
2. 蒸馏(Knowledge Distillation):用一个大型、高性能的教师模型去指导一个小型、量化或更高效的学生模型学习。学生模型不仅学习硬标签,还学习教师模型的“软目标”(logits分布),从而在保持性能的同时减小模型规模或适应低精度。
3. 微调(Fine-tuning)/增量学习(Incremental Learning):在少量目标领域的数据上对预训练模型进行再训练。这通常是最高效且最常用的方法,但需要少量有标注的目标领域数据。
4. 校准(Calibration):尤其对于量化模型,校准技术(如Post-Training Quantization, PTQ中的MinMax、Adel等)可以在量化后对激活值或权重范围进行调整,以最小化精度损失,避免需要重新训练。

这些方法都是为了让模型在面对与训练时不同的数据或计算环境时,仍能保持其鲁棒性和高性能。

问得好!对于小模型或者卡不够用的穷人……哦不,个人开发者来说,8位量化依然有它存在的意义,但核心收益可能从“吞吐量爆炸”变成了“省显存、省能耗”。你想啊,小模型本来就跑得快,它的瓶颈可能不是计算速度,而是怎么能塞进有限的显存里。量化能让模型体积变小,这样就能在更便宜的卡上跑,或者在同一张卡上多跑几个模型,这可是实打实的省钱!

除了量化,还有一些很给力的“黑科技”适合中小模型:
* 模型剪枝: 把模型里那些“不干活”或者“偷懒”的部分剪掉,模型立马精神百倍。
* 知识蒸馏: 找个超级大佬(大模型)当老师,让小模型跟着学,学到大佬的精髓,但自己又很“轻巧”,这是效率党的最爱!
* 用“瘦身款”模型: 比如MobileNet、ShuffleNet这些,它们就是为了在手机上、边缘设备上跑而设计的,天生就“瘦”。

所以,就算没有H100,我们也能玩转AI,只是玩法更精致、更抠门……啊不,更高效罢了!