吞吐提升76%!小红书开源RL训练引擎Relax
小红书AI平台团队刚刚开源了Relax——一个为全模态数据、Agentic工作流和大规模异步训练协同设计的现代RL训练引擎随着RL后训练逐步延伸至全模态与Agentic场景多模数据异构、系统稳定性和角色耦合等方面的问题日益凸显。为此小红书AI平台团队刚刚开源了Relax——一个为全模态数据、Agentic工作流和大规模异步训练协同设计的现代RL训练引擎实测全异步Off-Policy模式相比共卡On-Policy吞吐提升76%相比veRL的全异步实现提升20%先说痛点全模态Agentic RL训练的三重困境RL后训练正在经历两个根本性变化模型走向全模态图文音视频一起训甚至带有音频或图片视频输出训练走向Agentic多轮推理、工具调用、复杂环境交互。这两个趋势叠加把现有RL训练框架逼进了三重困境——而且这三个问题不是各自独立的它们环环相扣必须一起解决。困境一数据异构。高质量的图片和音视频原始数据传输体积大、CPU预处理开销高、编码后token爆炸multi-modal encoder无法和已有并行策略高效协同——在小红书内极其丰富的多模场景下需要一款定制优化的框架。困境二系统脆弱。多模态下较高的OOM风险叠加上千卡长时训练硬件故障、NCCL超时随时出现——传统方案缺乏分钟级故障恢复和单角色弹性伸缩能力。困境三角色耦合。Colocate方案下各角色共享GPU只能串行执行Trainer干等最慢的Rollout完成现有全异步方案虽然把Rollout和Train拆到不同组但缺少细粒度的流水线调度。核心洞察三者耦合必须协同设计这三个困境形成了一条因果链多模态让系统更低效且脆弱→催生服务化隔离和全异步架构→催生数据总线→总线的字段存储天然兼容多模态→三个问题一条因果链完成闭环。Relax用一套协同设计co-design把多模态原生、服务化容错、全异步流水线一并解决。全异步架构Rollout和Train并行一秒都不浪费核心思路拆开来并行跑传统方案把Rollout和Train绑在一起。Relax的做法很暴力——直接拆成两个独立服务。Rollout服务专门做推理生成用SGLang引擎Train服务专门做梯度更新用Megatron后端两个服务通过TransferQueue数据总线连接。Rollout生成的数据往里写Train从里面读互不阻塞。这带来一个直接的好处Rollout的推理时间大部分被掩盖Train不用等Rollout完事再开工而是上一批数据训着下一批数据就已经在生成了。Micro Batch级别的流水不等慢的那个光把Rollout和Train拆开还不够。传统方案——包括现有的全异步框架——还有一个致命缺陷全局batch同步。Relax把粒度从全局batch推进到了micro batch级别的流水线例如把256条的全局batch切成32条一组的micro batch每组生成完立即写入TransferQueue下游的Advantage、Train各阶段一看有数据了马上开始处理——不同micro batch在不同阶段形成流水线并行。效果长尾样本再也拖不死整个step。单条慢样本只影响它所在的micro batch其他micro batch正常流转对于超时仍未完成的样本Relax还支持partial rollout——已生成的部分直接回收利用不白等也不白扔。整条RL训练链路从粗粒度的“Rollout完了再Train”变成了细粒度的多阶段流水线。训练不怕挂服务化隔离与自动恢复服务化隔离一个挂了不拖累全局Relax不是一个铁板一块的单体程序。它把RL训练的每个角色——Actor、Ref、Rollout、Reward——拆成独立的Ray Serve服务。每个服务有自己的进程空间、故障域和资源配额。它们之间不直接引用全部通过TransferQueue数据总线通信。说白了——Ref挂了Actor照样训。一个Rollout实例OOM了其他Rollout正常工作。传统方案呢一个节点出问题整个任务重启进度全丢两级恢复策略能小修绝不大动Relax有一套完整的健康监控体系通过主动上报和心跳超时两条路并行并根据角色重要性分级处理简单来说——能局部修就局部修必须全局重启也从最近的checkpoint恢复而不是从头再来。分布式Checkpoint服务DCS快速恢复的关键恢复快不快核心取决于权重能不能快速到位。Relax专门为此搞了一个独立服务——DCS它不是简单的“训完了存一下”而是一个带拓扑感知的分布式权重传输系统。DCS自动发现TP/PP拓扑、构建跨异构并行的rank映射集群内走NCCL GPU通信最低延迟跨集群走TCP对弹性扩缩甚至是联邦集群下的弹性扩缩足够友好。弹性伸缩训练不停Rollout随时加减Relax支持训练过程中动态扩缩Rollout实例提供Ray原生集群内增减副本和外部引擎接入外部集群两种模式。每个新实例依次经过PENDING→CREATING→HEALTH_CHECKING→WEIGHT_SYNCING→READY→ACTIVE六个阶段每一步都有超时和回滚任何一步失败自动清理资源。缩容同样不粗暴——先停止接新请求、排空在途任务、等待权重更新完成再逐步下线。Relax还内置了基于KV Cache利用率、排队深度、TTFT等指标的自动伸缩器。多模不再拖后腿CPU预处理和ViT并行全解决Processor Pool。多模态processor图片resize、视频抽帧、音频重采样是CPU密集型操作单线程跑会被GIL卡死。Relax用ProcessPoolExecutor把processor放到独立进程共享内存零拷贝传输配合asyncio实现processor、media encoding、Rollout请求三阶段流水线——CPU预处理完全藏在GPU推理延迟背后。ViT原生HFTP维度数据并行。通过Megatron Bridge保持ViT的HF原生实现在TP维度上复制到所有rank各rank独立encode不同的pixel_values切片后AllReduce合并。ViT参数只占1-5%冗余开销可忽略省掉格式转换的工程负担和精度风险。异步通信。Relax基于TransferQueue的async_put / async_get让数据读写与GPU计算完全重叠配合micro batch流水线实现端到端的细粒度异步。不只是快Agentic RL也原生支持服务化架构带来的另一个好处——Agentic场景天然适配。自定义RolloutAgent要调工具、查数据库、跑沙箱写成可插拔服务挂上去就行。多轮状态管理服务化架构下多轮交互的状态管理就是服务间的消息传递。灵活的RewardRule-based、LLM-as-JudgeGenRM、自定义函数按需组合。实验结果数据说话实验环境NVIDIA H800 80GB GPU集群NVLink/InfiniBand互联。全异步吞吐提升76%在Qwen3-4BDAPO-MATH-17k16×H800的配置下全异步Off-Policy模式相对Colocate方案吞吐提升76%相比VeRL的全异步实现提速20%28.7 vs. 23.9 steps/hour即使是分离异步下的On-Policy也可提速12%。加速来自三层叠加1. 流式micro batch调度不等慢样本生成一组处理一组即使是async的on-policy也能拿到提速收益2. 资源分离log-prob和reference log-prob在独立GPU上计算计算量完全被掩盖。3. 无sleep/wakeup开销独立集群部署不再需要卸载/重载模型参数4B模型省10秒35B模型可能省50-80秒。收敛质量快了但没降质最关键的问题跑这么快效果打折了吗按wall-clock time看三种模式最终收敛到相同reward水平但Async Off-Policy达到同等reward的wall-clock时间比Colocate缩短43%。更快收敛效果不掉。近乎无损的R3MoE模型在RL训练中有个隐蔽的坑Rollout和Training阶段的expert路由可能不一致导致log probs mismatch从而加重off-policy。Relax的R3Rollout Routing Replay在Rollout时记录路由决策Training时原样回放Mismatch降低约38%Relax额外开销仅1.9%veRL为32%。为什么差这么多因为Relax的异步pipeline天然吸收了R3的开销——序列化路径重写异步device-to-host传输R3完全跑在关键路径之外。总结当RL训练从纯文本单轮走向全模态Agentic数据异构、系统脆弱、角色耦合三重困境不再是可以分别解决的独立问题。Relax的回答是一套协同设计全模态原生pipeline解决数据异构服务化隔离DCS快速恢复解决系统脆弱micro batch级全异步流水线解决资源利用率——三者因果闭环缺一不可。Relax的落地有赖于开源社区的支持研究团队希望在此致谢感谢Slime和SGLang团队在RayMegatronSGLang结合用于RL训练方面的基础性工作Relax正是在此基础架构之上构建的。感谢NVIDIA Megatron Bridge团队开发的开源checkpoint转换框架Relax将其扩展以支持全模态模型。感谢华为昇腾TransferQueue团队开发的开源异步数据总线它是Relax分布式架构的数据面核心。
本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.coloradmin.cn/o/2521695.html
如若内容造成侵权/违法违规/事实不符,请联系多彩编程网进行投诉反馈,一经查实,立即删除!