Molt: A Scalable PyTorch-Native Training Framework for Agentic Reinforcement Learning
Jian Hu, Huiying Li, Hao Zhang, Binfeng Xu, Yifan Zhang, Shaokun Zhang, Hemil Desai, Michael Demoret, Pavlo Molchanov, Jan Kautz, Yi Dong
cs.LG, cs.CL, cs.DC
2026-07-23
英伟达的 Molt 把 agent RL 训练框架压到 8.6K 行,单个异步循环训多模态和 MoE 策略、绝不训自己没生成的 token;在 Qwen3-30B-A3B 上与 Megatron 加 SGLang 栈吞吐统计持平,还跑通了 700B MoE。
agent 强化学习的研究是不断改算法:换估计器、加管线阶段、改 rollout 方式。可在主流框架里,每改一处都要穿过 trainer、分布式后端、rollout 胶水层一堆代码,改一次的代价全压在研究者身上。框架本身变得比要验证的假设还复杂。Molt 要解决的就是这个:把框架做小到研究者能整个装进脑子、AI 编程助手能从头到尾读懂,改算法的端到端成本尽量低,同时吞吐不能输给 SOTA 的重型栈。
Molt 的取舍很激进。全框架自有的 RL 代码约 8.6K 行(verl 约 62K、slime 约 25K),只用一个训练后端(AutoModel 加 FSDP2)和一个推理引擎(vLLM),都不 fork。四个核心概念:agent 是普通 Python 程序(Gym 风格环境,或用原生 OpenAI、Anthropic SDK 的 ChatAgent)、generator 做 token 精确捕获、trainer 只有一个可见的训练循环、estimator 和 loss 是纯函数。
几个硬性正确性不变量撑住它:token 同一性(采到的 token id 定义轨迹)、policy 版本语义(可训练 token 保留行为策略的对数概率)、前向一致性(rollout 和 actor 在模型语义上一致)。工程上靠 streaming pool 让 prompt 组常驻、按完成顺序吐训练 batch;partial rollout 能暂停引擎、广播 actor 分片、再恢复,不掉在途请求;token 进 token 出(TITO)全程不离开 token 空间,消灭重分词漂移;权重通过 NCCL 直接从 actor 广播到引擎。组合式并行:FSDP2 加 AutoModel 原生的 tensor、专家、上下文并行,扩规模是改配置不是迁移。
代码量上,Molt 完整 RL 路径约 8.6K 行,verl 约 62K,slime 约 25K。引擎特性实测(Qwen3.6-35B-A3B 多模态 MoE,2 节点 8 训加 8 推):投机解码配 MTP 头,单步生成从 329s 降到 64s(5 倍);optimizer CPU offload 把 actor 峰值显存从 64.7GB 降到 46.4GB,policytrain 时间加 18%。
正面硬刚:在 Qwen3-30B-A3B、DAPO-Math 上,完全异步分离式 8 加 8 协议对比 slime(Megatron-Core 加 SGLang)。
| 配置 | 单步 (s) | Tok/GPU/s |
| Molt (AutoModel 加 vLLM) | 119.4 ± 2.3 | 461 |
| slime (Megatron 加 SGLang) | 109.5 ± 10.3 | 502 |
作者明确不主张任何一方更强:slime 跨 run 的步时方差(102 到 121s)和 Molt 的带(119.4 ± 2.3)重叠,均值差约 9% 落在波动里。规模上演示了 700B MoE、专家并行 256 跑通完整异步循环。
对做 agent RL 的研究者,这是「框架不该成为研究瓶颈」的一份证据:用不到别人七分之一的代码,在吞吐上打平 Megatron 重型栈,而且代码量小到人和 AI 编程助手都能完整读懂、端到端改。它也示范一种工程判断,扩规模靠组合已有并行(FSDP2 加 AutoModel 原生并行加 vLLM),不用重造分布式后端。开源(NVIDIA-NeMo/labs-molt)带 recipe 和容器。
要诚实:它主打研究场景,不提供企业级数据转换层和控制面;吞吐对比是「统计持平」不是「更快」,而且 30B 检查点暴露了上游分布式 MoE 前向不一致的问题,只做了吞吐对比、收敛一致性待上游修复后再验。
作者自述:部署选择更窄(故意为单后端让路);部分组合(如 packed batch 配上下文并行)在配置期就被拒;不主张吞吐优越;基准受限。
存疑点:「打平 Megatron」是单点协议(一个模型、一个数据集、8 加 8 卡)下的吞吐持平,泛化到别的模型规模和任务未验证;700B MoE 只演示「跑通」,没给收敛或质量数据。作为极简框架,它把复杂度外包给了 vLLM、Ray、AutoModel 这些依赖,这些依赖的版本和 bug 风险,是 8.6K 行省下来的代价的另一面。