HyperAIHyperAI

Command Palette

Search for a command to run...

Molt:面向智能体强化学习的可扩展 PyTorch 原生训练框架

摘要

智能体强化学习研究涉及持续的算法修改、新的估计器、新的流水线阶段和新的 rollout 方案,而在主流框架中,每一次变更都会贯穿训练器、分布式后端和 rollout 粘合代码的多个层次:其成本在每次迭代中都落在研究者身上。Molt 是一个 PyTorch 原生的训练框架,旨在将这一成本保持在较低水平:其代码库紧凑且清晰,足以让研究者了然于胸,也能让 AI 编程助手完整阅读和推理,从而使算法流程能够被端到端地追踪和修改。智能体是一个普通的程序,一个异步循环在训练多模态和混合专家策略的同时,绝不训练任何非自身生成的 token,在 token、策略版本和模型语义上保持一致。精简并不以性能为代价:在匹配的全异步协议下,Molt 在统计上与基于 Megatron 的最先进技术栈相当。Molt 是开源的,并在 https://github.com/NVIDIA-NeMo/labs-molt 提供配方和容器。

一句话总结

NVIDIA研究人员推出了Molt,一个紧凑的PyTorch原生训练框架,用于agent驱动强化学习。它采用异步循环训练多模态和混合专家策略,同时确保模型绝不在未生成的token上进行训练,实现了与最先进的基于Megatron的系统具有统计可比性,并提供了一个足够紧凑的代码库,可供人类完全掌握,也可供AI编码助手完整阅读和推理。

核心贡献

  • Molt提供了一个用于agent驱动强化学习的PyTorch原生训练框架,其紧凑的代码库旨在供研究人员或AI编码助手完整阅读和推理。
  • 该框架通过仅使用自生成的token训练多模态和混合专家策略、保留行为对数概率以及使用路由器重放来防止因推理和训练路由不一致导致的MoE不稳定性,从而强制实现token级别的一致性。
  • 在匹配的全异步协议下,Molt实现了与最先进的基于Megatron的堆栈的统计同等性,并将相同的精简循环从4B扩展到700B MoE参数,专家并行度达到256。同时,开源版本包含了配方、容器以及一键复现结果的方法。

引言

作者们针对agent驱动强化学习中日益加剧的矛盾:主流RL框架专为大规模训练工作负载而设计,施加了分布式后端、控制器和配置的多层复杂性,从而拖慢了研究中至关重要的快速算法修改。原型化一个新想法常常需要在众多胶水层中追踪更改,而服务与训练之间细微的数值不匹配(如token化或专家路由的分歧)可能会悄无声息地偏置梯度而不会引发错误。

为了减少这种摩擦,作者们推出了Molt,一个精简的PyTorch原生框架,将人类和AI编码助手的可读性作为首要设计约束。其主要贡献是一个代码库比可比基于Megatron的生产堆栈小数倍的框架,在匹配吞吐量的同时,强制了三个核心正确性不变条件:仅使用生成的token ID进行训练、保留行为策略的对数概率,并确保推理执行者的前向一致性。

方法

作者们将Molt构建为一个精简的agent驱动强化学习系统,特意避开了大规模训练堆栈典型的分层间接方式。整个运行时围绕一个异步循环组织,连接三个定义明确的组件组:一个agent池、一组vLLM推理引擎以及一个基于NVIDIA AutoModel并采用全分片数据并行2(FSDP2)、专家并行(EP)和上下文并行(CP)的可训练策略Actor。基于Ray的队列将生成过程与训练解耦,保持prompt组处于运行状态,并在收集到足够完成结果后立即输出训练批次。该架构避免了任何适配层或参数服务器,而系统绝不在未生成的token上进行训练这一不变条件则通过token优先契约来强制实施:每条消息都携带token ID、每个token的对数概率以及对齐的多模态张量,并且没有任何组件会重新对文本进行token化。

agent边界遵循算法自身对象的模块化原则。研究人员提供一个AgentRunner模块,其余部分为普通Python代码。在Env模式下,框架驱动LLM循环,在每次模型动作后调用用户的step函数,并处理token化和预算。在ChatAgent模式下,使用OpenAI或Anthropic线路协议的现有agent代码可以原样使用:Molt启动一个环回聊天服务器,将每个请求解码为token精确的累积,从而在不编写任何集成代码的情况下实现token输入/token输出捕获。上下文压缩(摘要化或丢弃之前的轮次)通过分段轨迹自动处理,并保留推理结束时的奖励。两种agent形式使用相同的聊天模板数据,且单个数据集可服务于任一路径,而不会造成数据路径混淆。

传输过程通过构造保持正确。系统不包含任何vLLM分叉;所有交互都利用稳定、文档化的端点。面向训练器的生成过程使用引擎的token级接口,将prompt和补全结果作为token ID以及每个token的对数概率。请求路由器将同一推理的所有请求固定到同一引擎,以保留前缀缓存,而多模态prompt则由服务端渲染,以避免静默图像丢失。权重更新通过从actor直接到每个引擎的NCCL广播进行处理,完全绕过路由器。

规模通过配置而非迁移来表达。用于训练4B密集模型的启动脚本同样通过简单的并行标志(如--fsdp.ep_size 256)表达大型混合专家(MoE)配置。对于MoE架构,采用本地推理路由重放:引擎返回其每个token的专家选择,而actor在训练期间重放这些选择,从而消除因推理和训练做出不同路由决策而导致的故障模式。其他优化(如推测性解码、前缀缓存和CUDA图)以引擎标志的形式出现,而将优化器卸载到主机内存以及保守的防护栏在启动时拒绝不兼容的并行组合。

算法层同样避免了继承层次结构。优势估计器是奖励和组别的纯函数,通过名称选择(例如,带组平均基线的REINFORCE++、RLOO、GRPO、带评论家的GAE或在策略蒸馏)。全局损失归一化使用整个批次的token均值,使更新对数据并行大小和梯度累积保持不变。带有序列级门控的离策略修正,结合动态批次过滤,处理异步推理,并确保每个token的训练对数概率与生成时的值保持一致。最终结果是,算法上的改变(例如引入一个新的优势估计器)只需触及单个函数、一个标志以及一个记录指标,而整个循环在一次通读中即可理解。

实验

评估通过测量RL堆栈的内存占用、引擎优化对吞吐量的影响以及其与基于Megatron基线相比的性能,量化了精简化的代价。该框架仅包含8600行代码,而上游引擎特性(前缀缓存、推测性解码、优化器卸载)以配置标志的形式出现,带来了明显的加速和内存节省。在30B MoE模型上进行的受控一对一比较显示出统计上可比的每步时间(119.4 ± 2.3秒对109.5 ± 10.3秒),证实了这种组合式、精简的设计在不增加性能开销的情况下实现了与完整堆栈的吞吐量平等。

在非聚合GPU上的匹配异步协议下,Molt和slime实现了统计上可比的每优化器步时间,运行间变异性重叠,没有明显胜者。训练布局是一个一阶因素:不匹配的并行设置在这项16K工作量上可能导致步时间增加大约30%。基准测试端的MoE前向不匹配导致128个专家的检查点无法进行有效的策略更新,因此报告的步时间仅衡量吞吐量;而35B的工作量则没有此间隙。在受控协议下,Molt和slime的步时间在运行间变异性范围内重叠,表明后端性能没有显著差异。将为32K上下文优化的上下文并行布局强制用于16K工作量,会使Molt的步时间减慢约30%。基准检查点中的分布式MoE对数概率不匹配导致序列门控拒绝批次,因此优化器步仅衡量吞吐量而无有效策略更新。随着输出长度向推理或agent领域增长,生成过程主导步时间,训练后端的相对贡献缩小至无关紧要。

在Qwen3-30B-A3B模型和16K上下文下,精简的Molt堆栈与基于Megatron的slime堆栈的直接比较显示统计上可比的优化器步时间(119.4 ± 2.3秒对比109.5 ± 10.3秒)。约9%的平均差异落在slime的跨运行变异性范围内,表明两者均无吞吐量优势。在较长的输出长度下,训练后端在总步时间中的占比缩小,进一步强调了实际等同性。当Molt(FSDP2 + vLLM)和slime(Megatron-Core TP4+SP + SGLang)在8+8块H100 GPU上使用各自推荐的并行布局时,它们提供了可比的每步吞吐量。slime的步时间分布为102–121秒,完全覆盖了Molt的117–122秒区间,使差异在统计上不显著。两个堆栈均在非聚合推理和训练下完全异步运行,且均未加载参考模型,从而确保每步的算法工作完全相同。基准检查点中的分布式MoE前向不匹配导致actor对数概率偏差约1 nat,因此报告的步时间仅反映吞吐量而无有效策略更新。

在非聚合推理和训练的匹配协议下,两个堆栈实现了吞吐量平等:MOLT的步时间为119.4 ± 2.3秒,slime为109.5 ± 10.3秒,平均差异约9%,落在跨运行变异性范围内。由于两个堆栈都将生成与训练重叠,端到端步时间是有意义的吞吐量指标,而任何残余的后端优势会随着生成长度的增加而缩小。报告的每GPU每秒token数仅反映生成吞吐量,因为检查点不匹配阻止了有效的策略更新。步时间分布在不同运行中重叠,表明没有任何堆栈在速度上具有统计意义上的优势。slime产生502 tokens/GPU/s,而MOLT产生461 tokens/GPU/s,但slime运行时间的分布(102–121秒)涵盖了MOLT的平均值。在非聚合GPU上,生成和训练完全重叠,因此步时间是端到端吞吐量指标。9%的平均步时间差异处于跨运行变异性内,因此没有一个堆栈声称优越。随着输出长度增长至推理级轨迹,后端特有的开销相对于生成时间变得可以忽略。16K上下文长度是最不利于隐藏后端差异的区域,然而两个堆栈已经具有可比性。将为32K上下文调优的上下文并行布局强制用于此序列长度,会使MOLT的步时间增加约30%。分布式MoE检查点不匹配导致actor对数概率偏差约1 nat,使得这些测量结果仅为吞吐量而无有效策略更新。

在非聚合GPU上的匹配异步协议下,对Molt(FSDP2 + vLLM)和slime(Megatron-Core与SGLang)堆栈进行直接比较,显示统计上无法区分的步时间,运行间变异性重叠,没有任何后端获得有意义的速度优势。训练并行布局是一个一阶因素:将为32K上下文调优的上下文并行配置应用于16K工作量,会使步时间增加约30%。在128个专家的基准测试中,分布式MoE前向不匹配阻止了有效的策略更新,因此报告的步时间仅衡量吞吐量;而35B的工作量则没有该间隙。随着输出长度向推理或agent领域增长,生成过程主导总步时间,训练后端的相对贡献缩小至实际无关,再次确认了两个堆栈的等同性。


用 AI 构建 AI

从创意到上线——通过免费 AI 协同编码、开箱即用的环境和最优惠的 GPU 价格,加速您的 AI 开发。

AI 协同编码
开箱即用的 GPU
最优定价

HyperAI Newsletters

订阅我们的最新资讯
我们会在北京时间 每周一的上午九点 向您的邮箱投递本周内的最新更新
邮件发送服务由 MailChimp 提供