Forge: Scalable Agent RL Framework and Algorithm
Forge: Scalable Agent RL Framework and Algorithm
MiniMax
1. 主要贡献
本文介绍了MiniMax为解决大规模强化学习(RL)在复杂真实世界Agent应用中面临的“不可能三角”(系统吞吐量、训练稳定性和Agent灵活性)而开发的内部RL框架——Forge。主要贡献如下:
1. 整体架构创新:提出了Forge框架,通过灵活的系统架构、标准化的交互协议(Middleware设计),解耦了Agent推理逻辑与底层训练设施,支持任意Agent脚手架(Scaffold)的训练。
2. 大规模工程实践:该系统支持了MiniMax M2.5模型的构建,处理了超过10万种不同的真实世界Agent脚手架和环境,支持长达200k的上下文长度,保持了每日百万级样本的处理吞吐量。
3. 算法与调度优化:引入了Windowed FIFO混合调度策略解决效率与分布一致性的冲突;提出了前缀树合并(Prefix Tree Merging)技术消除计算冗余;采用CISPO算法结合混合域训练和复合奖励框架(包含时间成本奖励),实现了高效且稳定的训练。
2. 问题定义与背景
在详细介绍架构之前,论文首先将Agent RL系统的优化目标形式化为最大化有效Agent训练产出($J$):
$$ \max J(\theta) = \text{Throughput}(A) \times \text{Sample Efficiency}(A) $$ $$ \text{s.t. } \forall A \in \Omega_{agent} (\text{Arbitrary Agent}) $$其中系统吞吐量受限于Rollout、训练、数据处理和I/O四个组件;样本效率取决于数据分布、质量、算法效率和Off-policy程度。
1.1 Agent扩展性与框架灵活性挑战
当前RL范式存在“玻璃天花板”:
* 受限的Agent自主性:标准框架将Agent视为白盒函数并共享状态,难以模拟复杂的认知架构(如动态上下文管理、多智能体协作),阻碍了模型在任意黑盒Agent上的泛化。
* Token一致性障碍:现有的TITO(Token-In-Token-Out)架构强制Agent与底层Token逻辑耦合。在复杂的上下文管理(Context Management, CM)下,维持高层推理抽象与底层训练表示的一致性计算成本极高。
1.2 系统效率与计算冗余
Agent的Rollout完成时间方差极大(从秒级到数小时),导致调度死锁:
* 异步控制器困境:严格的FIFO/同步调度受“掉队者效应”(Straggler Effect)影响导致集群空闲;而贪婪/FFFO模式虽然最大化了吞吐量,但会导致严重的数据分布偏移(Data Distribution Shift),造成非平稳训练环境和梯度震荡。
* 前缀冗余:Tokenizer机制与上下文管理的相互作用导致大量请求共享相同的前缀,在训练期间造成显著的计算浪费。
1.3 算法挑战:信用分配与优化稳定性
* 稀疏奖励与高梯度方差:Agent任务涉及超长视界和延迟反馈,在200k上下文窗口内将信用分配给特定Token极其困难,导致低信噪比和高梯度方差。
* 忽略延迟的优化:传统RL目标仅关注正确性而忽略执行的时间成本。在现实场景中,不同轨迹的延迟差异巨大,标准范式未能激励并行化或高效工具的使用。
3. 方法细节
2. 系统架构与Agent RL范式
为了缓解“效率 vs. Off-Policyness”的权衡并减少冗余,系统采用了以下架构创新。
2.1 RL系统设计
为了实现真正的可扩展架构,系统超越了具体实现,采用了通用的“中间件”设计,将Agent的推理逻辑与底层训练基础设施解耦。RL系统由以下3个模块组成:
-
Agent Side(Agent端):
- 通用抽象:该层对通用Agent(包括白盒和黑盒架构)及其运行环境进行抽象。
- 纯轨迹生产者:它编排递归的环境交互,使Agent作为纯粹的轨迹生产者运行。通过将环境反馈与系统开销解耦,Agent可以专注于核心业务逻辑(如上下文管理和推理链),而无需感知底层的训练和推理机制。
-
Middleware Abstraction Layer(中间件抽象层):
- 物理隔离:作为桥梁,该层在物理上隔离了Agent端与训练/推理端。包含Gateway Server和Data Pool。
- Gateway Server(网关服务器):作为一个标准化的通信网关,处理Agent与LLM之间的完成请求。通过使用通用标准协议,有效隔离了底层模型的复杂性与Agent的高层行为逻辑。
- Data Pool(数据池):作为分布式数据存储,异步收集来自Agent的rollout轨迹和报告。它充当缓冲区,解耦生成与训练,允许用户应用灵活的数据处理和批处理策略,以提高训练效率和算法利用率。
-
Training and Inference Side(训练与推理端):
- 计算承载:负责繁重的计算任务,由LLM Engine和Train Engine组成。
- Rollout Engine:致力于高吞吐量的Token生成,响应中间件转发的请求。
- Train Engine:从Data Pool消费处理后的Token序列以更新策略。它保持与LLM Engine的同步,确保Agent使用最新的策略分布进行探索。
通过这种模块化设计,Forge实现了引擎与Agent的完全解耦,集成了数百种类型的脚手架和数千种不同的工具调用格式。
2.2 针对上下文管理(CM)的白盒Agent RL
对于白盒Agent,通过全面的脚手架设计和增强,直接观察并优化模型在特定架构上的性能。针对长视界任务中上下文管理(如DeepSearch)面临的问题,采取了以下措施:
- 解决Context Rot(上下文腐烂)与不匹配问题:
- 问题分析:交互轮次增加导致中间推理步骤积累,产生“注意力稀释”效应;仅在推理时应用CM会导致与RL训练数据的严重分布偏移(Inference-Training Mismatch)。
- CM驱动的状态转换:将CM建模为显式的Agent动作,上下文转换自然嵌入环境动态中。从 $S_t$ 到 $S_{t+1}$ 的状态转换隐式封装了上下文切换逻辑,将上下文适应直接折叠进模型的训练目标中。
- 自适应推理模式:通过在此框架内优化策略 $\pi$,模型学会了内化分布偏移,涌现出优先考虑“状态关键”Token的鲁棒推理模式。
- 上下文感知策略:模型被训练为在RL生成过程中预测潜在的CM操作,主动保留任务关键信息并修剪无关噪声。
2.3 黑盒Agent RL:异构脚手架的鲁棒性
针对用户操作的专有或复杂“黑盒”Agent架构,Forge确保了跨不同认知架构的一致优化。
- 非侵入式集成与兼容性:Forge对Agent的内部实现细节保持完全不可知。Agent只需将请求路由到RL服务Gateway,框架自动处理数据收集和训练。因此,Forge无缝支持任意上下文操作(如记忆压缩、历史重写)以及复杂的内部Agent循环(如Deep Think、Multi-Agent架构)。
- 多脚手架泛化:通过将训练循环与Agent内部状态解耦,M2.5实现了与大量黑盒Agent的广泛兼容性,覆盖了从严重依赖Sandbox和模型上下文协议(MCP)的代码中心Agent(如将OpenCode Agent作为黑盒训练),到采用激进上下文缩减策略(如Truncate BC)的Agent。实验结果表明,这种方法在完全不透明的黑盒系统上也能提供一致且稳定的改进。
3. 工程优化
3.1 混合调度策略:Windowed FIFO
为了解决系统吞吐量与分布一致性之间的冲突,引入了Windowed FIFO策略。该策略在训练调度器上施加滑动约束,作为严格同步排序与贪婪异步执行之间的“中间地带”。
-
受限的可见性范围(Restricted Visibility Scope):
- 设生成队列为 $Q = [T_0, T_1, ..., T_{N-1}]$,当前头部索引为 $i$。训练调度器被严格限制只能从 $[T_i, T_{i+W-1}]$ 范围内提取已完成的轨迹,其中 $W$ 为窗口大小(例如 $W = 0.3N$)。
-
窗口内的局部“贪婪”无序(Local "Greedy" Disorder):
- 在活动窗口 $[T_i, T_{i+W-1}]$ 内,调度器可以立即检索任何已完成的轨迹。这减轻了队头阻塞(HoL)效应,因为窗口内的快速任务无需等待绝对的第一任务完成。
-
窗口边界的全局“严格”阻塞(Global "Strict" Blocking):
- 即使索引 $j > i+W$(窗口外)的任务已完成(这在大型生成批次中的简单快速任务中很常见),调度器也被禁止获取它。
-
约束实现:
- 只有当头部的任务被消费时,窗口才会向前滑动($i \to i+1$)。这种机制强制调度器等待当前窗口内的“掉队者”(复杂、长视界任务),防止训练分布向队列后方“快速且简单”的样本漂移。
3.2 基于前缀树合并(Prefix Tree Merging)的Agent轨迹训练加速
在Agent训练中,数据集通常由大量的多轮对话样本组成,这些样本在结构上表现出高度的重叠。
- 传统方法的冗余挑战:
- 前缀重叠:在简单的多轮对话中,消息被顺序追加。理论上,共享相同历史的多个补全可以合并。
- 复杂的上下文管理:Agent经常采用复杂的CM策略,导致不同的补全经常共享广泛的公共前缀。
- 朴素方法的局限:传统方法将每个样本视为独立实体,反复重新计算这些公共前缀。在长上下文场景中,这导致了TFLOPS的大量浪费。
- 前缀树合并方案:
- 树结构转换:将训练过程从“线性处理”转换为“树结构”方法。
- 合并逻辑:针对Agent场景中的复杂上下文管理,只要多个补全共享底层前缀,即使后续响应略有不同或属于不同的采样分支,也可以在样本级别合并为单个前缀树。
- 技术实现:利用注意力原语(如Magi Attention),确保逻辑执行与标准前向传递保持一致。前向传递后,根据元数据解构前缀树以正常计算损失,确保对下游逻辑零影响。
- 效果:通过消除冗余的前缀预填充(Prefilling),该方案实现了40倍的训练加速,并显著降低了内存开销以支持更长的序列或更大的批量大小,同时保证了与标准方法在数学上的严格等价性。
3.3 极致的推理加速
通过三项架构创新优化生成管道:
- 基于MTP的投机解码(Speculative Decoding):不使用静态草稿模型,而是使用通过Top-K KL损失持续微调的多Token预测(MTP)头。这确保了与不断演变的RL策略对齐,通过减轻分布偏移维持高接受率和显著加速。
- 异构PD分离(Heterogeneous PD Disaggregation):解耦预填充(Prefill)和解码(Decode),以消除混合MoE调度中的PD干扰,并允许每个实例采用独立的并行策略,同时最大化全局吞吐量并优化长视界任务的尾部延迟。
- 全局L3 KV缓存池:为了防止多轮Agent RL中的冗余预填充并通过组级rollout最大化前缀缓存命中率,引入了基于DFS的全局L3缓存。成本感知调度器通过权衡排队延迟与缓存迁移成本来动态路由请求,在不过载实例的情况下最大化缓存局部性。
4. 可扩展Agent RL算法
4.1 RL算法
核心算法采用CISPO,并专门针对长视界Agent的特性进行了调整。
* 统一混合域训练:不同于导致负迁移或域间干扰的多阶段强化学习,采用统一训练策略。同时混合推理(Reasoning)、通用问答(General QA)和Agent领域的任务。这种联合训练方法减轻了顺序训练中常见的性能下降,并显著增强了模型在不同任务间的泛化能力。
提出了一个复合奖励框架,旨在解决超长上下文(高达200k)的信用分配挑战,同时确保训练稳定性:
* 过程奖励(Process Reward):为了提供密集反馈,针对中间行为(例如惩罚语言混合或特定工具调用错误),而不是仅依赖最终结果。
* 任务完成时间奖励(Task Completion Time Reward):在Agent场景中,将相对完成时间纳入奖励信号。这不仅取决于Token生成,还取决于特定工具执行和子Agent调用的延迟。这激励Agent利用并行性,从而加速任务执行。
* Reward-to-go用于方差减少:采用Reward-to-go公式对回报进行归一化,有效降低梯度方差并提高信用分配的精度,稳定优化过程。
4. 实验环境与结果
实验配置与环境
* 模型:MiniMax M2.5。
* 数据规模与上下文:支持高达200k的上下文长度,处理了超过10万种不同的真实世界Agent脚手架。
* 硬件与吞吐量:系统保持了每日百万级样本的处理吞吐量。
* 验证对象:包括OpenCode Agent(代码中心)、Truncate BC策略等多种黑盒与白盒Agent。
实验结果
* 训练稳定性与收敛性:通过CISPO算法与复合奖励框架,系统实现了奖励的持续收敛(见图2)。Reward-to-go机制有效降低了梯度方差。
* 训练效率提升:Prefix Tree Merging技术消除了前缀冗余,实现了40倍的训练速度提升,并显著减少了内存开销。
* 泛化能力:通过解耦架构和混合域训练,M2.5在数百种脚手架和数千种工具格式上展示了广泛的兼容性和性能提升,证明了在黑盒Agent上的鲁棒性。
5. 结论
MiniMax通过Forge框架成功解决了扩展Agent RL面临的“不可能三角”。通过整合灵活的解耦架构、Windowed FIFO调度、Prefix Tree Merging以及稳定的CISPO算法,Forge实现了RL系统吞吐量的突破,同时确保了跨任意Agent脚手架的鲁棒泛化。这一整体方法支持了MiniMax M2.5的大规模训练,交付了高效的真实世界Agent能力,推进了“Intelligence with Everyone”的使命。
💬 评论讨论
欢迎在这里分享您的想法和见解!