NVIDIA发布Molt框架简化Agent强化学习研究,8.6K行代码整合Ray/vLLM/NeMo,性能媲美Megatron,大幅降低Agent RL模型训练成本。
Agentic reinforcement learning 研究始终伴随着算法调整:新的 estimator、新的 pipeline 阶段、新的 rollout 方案。在主流框架中,每一次改动都要贯穿 trainer、distributed backend 和 rollout glue 等多层代码。研究人员每迭代一次,都要承担这样的成本。
NVIDIA NeMo 团队推出的 Molt,正是为了直接降低这种成本。它是一个 PyTorch-native 的 Agentic RL 框架,其设计目标颇为独特:代码库应该足够精简,让研究人员能够在脑中建立完整认知,也让 AI 编程助手可以通读整个代码库并进行推理。根据其公布的数据,从各框架的 RL 入口点出发追踪 import graph,Molt 的 RL 代码规模约为 8,600 行。采用同样的统计方法,verl 约为 62,000 行,slime 约为 25,000 行,OpenRLHF 约为 7,200 行。
没错,Molt 采用 Apache 2.0 许可证发布,同时提供了 launch code、Slurm 脚本和预构建容器。不过,研究论文将它定位为研究基础设施,而不是生产级训练服务;真正的门槛在于硬件。官方提供的 recipe 默认使用 2 个节点、每个节点配备 8 张 H100 GPU,其中 8 张用于训练,另外 8 张用于 rollout。
这意味着 Molt 主要适用于前沿或接近前沿水平的实验室、资金充足且从事 post-training 的 AI 初创公司、金融、医疗和机器人领域中使用私有环境训练 Agent 的企业 AI 研究团队,以及能够使用多节点 H100/H200 集群的高校实验室。其应用场景包括多轮工具调用 Agent、代码执行 Agent、vision-language 环境(官方提供了 geo3k recipe)、LLM-as-judge 奖励循环,以及将能力 on-policy distillation 到规模更小的 student model。
Molt 将 Ray 用于 placement 和 asynchronous queue,将 vLLM 用于 rollout,并使用搭配 FSDP2 的 NVIDIA AutoModel 进行训练。这三个组件均未被 fork,因此上游改进可以通过更新容器中的版本 pin 直接引入,无须进行 rebase。
它的 runtime 由一个 Agent pool、一组位于 request router 后方的 vLLM engine,以及一个可训练的 policy actor 构成。streaming pool 会持续维持正在处理的 prompt group,确保 actor 训练期间 engine 不会因请求耗尽而闲置。执行 partial rollout 时,系统会暂停 engine,通过 NCCL 将 actor shard 直接广播到每个 engine,然后恢复之前保留的请求,而不是将其丢弃。
一次 RL 运行只需指定一个导出 AgentRunner 的 Python module。除此之外,包括 reward 在内的所有部分都只是普通代码。
Molt 支持两种形式。使用 Env 时,框架通过与 Gymnasium 对齐的 step() 接管 LLM 循环。使用 ChatAgent 时,用户则可以通过标准 OpenAI 或 Anthropic SDK 自行控制循环。Molt 会启动一个同时支持这两种 wire protocol 的 loopback server,每个请求都会在服务端解码,并精确地按 token 累积到同一条轨迹中。当 long-horizon Agent 压缩 context 并重写 prefix 时,服务器会自动封存当前 segment,并开启一个新的 segment。
Molt 的设计围绕三条正确性 invariant 展开。第一是 token identity:定义 trajectory 的是采样得到的 token ID,而不是对 transcript 重新 tokenization 后的结果。第二是 policy-version semantics:可训练 token 会保留其 behavior policy 的 log-probability;对于异步场景,则在 sequence-level gate 之后逐 token 进行修正。第三是 forward consistency:rollout 与 actor 必须对模型语义保持一致。
对于 mixture-of-experts policy,最后一条 invariant 尤为重要。rollout router 和 training router 会各自独立选择 expert,而微小的数值差异就可能改变 top-k 选择结果。Molt 采用 rollout routing replay:由 vLLM 返回每个 token 对应的 expert ID,再由训练阶段的 forward 过程重放这些路由选择。
Molt 是一个采用 Apache 2.0 许可证的 Agentic RL 框架,RL 代码量约为 8,600 行,大约比 verl 小 7 倍。
Ray、vLLM 和 NVIDIA AutoModel 通过组合方式协同工作,均未被 fork,因此只需更新容器中的版本 pin,就能引入上游的新版本。
Agent 就是普通 Python 程序;通过一个 token-exact 的 loopback server,标准 OpenAI 或 Anthropic SDK 无须修改即可参与训练。
其吞吐量在统计意义上与基于 Megatron 的技术栈相当,同时官方明确披露了 MoE mismatch 这一注意事项。
扩展规模只需调整一个 flag:同一套循环既能运行 dense 4B 模型,也能通过 --fsdp.ep_size 256 运行 700B MoE 模型。
可以在这里查看论文和代码仓库。也欢迎在 Twitter 上关注我们,别忘了加入拥有超过 15 万成员的 ML SubReddit,并订阅我们的 Newsletter。等等!你使用 Telegram 吗?现在也可以加入我们的 Telegram 社群。
需要与我们合作,推广你的 GitHub Repo、Hugging Face Page、Product Release 或 Webinar 等内容吗?欢迎联系我们。
Asif Razzaq 是 Marktechpost Media Inc. 的 CEO。作为一位富有远见的创业者和工程师,Asif 致力于发掘 Artificial Intelligence 的潜力,用它推动社会进步。他最近创办了 Artificial Intelligence 媒体平台 Marktechpost。该平台凭借对 machine learning 和 deep learning 新闻深入、技术严谨且易于广大读者理解的报道而脱颖而出。平台每月浏览量超过 200 万,足见其在读者群体中的受欢迎程度。