Molt：面向智能体强化学习的可扩展 PyTorch 原生训练框架 Molt: A Scalable PyTorch-Native Training Framework for Agentic Reinforcement Learning
用约8.6K行可读代码搭出匹配Megatron栈吞吐的智能体RL训练框架
前置知识
FSDP2（Fully Sharded Data Parallel v2）
PyTorch 原生的全分片数据并行机制，将模型参数、梯度、优化器状态分散到多个 GPU 上，每个 rank 只持有完整参数的一个分片，前向/反向时按需 all-gather 还原。FSDP2 是新版实现，可与原生 tensor/expert/context 并行组合，是 Molt 中训练 actor 的基础并行手段。
Molt 不依赖 Megatron 的 TP/PP 抽象，而用 FSDP2 + AutoModel 原生并行来表达从 4B 稠密模型到 1T 级 MoE 的全部规模，理解 FSDP2 才能看懂为什么'规模是配置而非迁移'。
vLLM 与 PagedAttention
vLLM 是高性能 LLM 推理/服务引擎，用 PagedAttention 管理显存、提供前缀缓存、CUDA graphs、投机解码等优化。在 RL 训练中它常被复用为 rollout（采样）引擎，负责用当前策略权重生成轨迹。
Molt 把 vLLM 作为唯一的 rollout 引擎且不 fork 它，靠其 token 级接口做'逐 token 精确'采集。本文大量加速数字（投机解码 5×、前缀缓存 0.05s 重预填充）都来自 vLLM 的能力按 flag 直接流入。
GRPO / REINFORCE++ 等优势估计器
这是 LLM RL 中常用的无 critic（或带 critic）的策略梯度目标族。GRPO（Group Relative Policy Optimization）用同一 prompt 的多条采样做组内基线；REINFORCE++ 在此基础上加全局优势归一化稳定训练；RLOO 用留一基线；GAE 则配合 PPO critic。
Molt 的算法层把这些估计器实现成'纯函数'，理解它们的区别才能看懂作者为何强调'换一个优势估计器只是一个下午的改动、一个 flag、一个函数'。
MoE（Mixture of Experts）与路由重放
MoE 模型每层有多个专家网络，由路由器（router）按 token 动态选择少数专家计算。问题是训练侧和推理（rollout）侧各自独立做路由决策，微小数值差异会让两边选中不同专家、得到不同的稀疏计算图，造成'训练-推理不一致'。路由重放（routing replay）即把 rollout 时的专家选择记录下来，在训练时强行复用。
Molt 把 MoE 路由一致性列为三大正确性不变量之一，并用 R3 式路由重放在本地修复；这正是 35B MoE 基准出现约 1 nat 对数似然偏差、序列门拒绝整批的根本原因，是论文结果的关键细节。
TITO（Token-In / Token-Out）轨迹采集
传统的做法是先让模型生成文本、再训练前重新 tokenize，这会引入 token 序列漂移。TITO 要求全程不离开 token 空间：prompt 以 token id 输入、补全以 token id 加每 token 对数概率返回，训练直接消费这些原始 token，从根本上保证'训练的 token 就是被生成的 token'。
这是 Molt 的核心正确性保证，也是 ChatAgent 通过回环服务器让现成 OpenAI/Anthropic SDK 代码无需集成即可训练的原理，理解它才能理解作者的 token-first 契约。
研究动机
主流的 LLM 强化学习训练框架（如 verl、slime、NeMo-RL）是为超大规模训练而设计的，它们采用多后端结构：分离的 rollout 引擎、分布式 trainer、控制器、注册表、配置层。这种分层在超大规模场景下是合理的工程，但落到研究迭代上却成了沉重负担——因为智能体 RL 研究本质上是'持续的算法修改'：换一个优势估计器、加一个经验过滤阶段、改一种 rollout 方案，每一次改动都要穿过 trainer、分布式后端、rollout 粘合代码好几层才能生效，成本最终压在研究者每一次迭代上。更棘手的是，智能体在线 RL 有'安静的失败模式'：服务引擎与 actor 名义上评估同一个策略，但 tokenization、采样变换、多模态渲染、权重版本、MoE 路由的细微差异不会报错，只会让梯度悄悄偏置或被拒绝；现有 harness 自己掌管上下文/工具/控制流，而 trainer 又通常期望框架特定的环境接口，两者之间存在一条难以弥合的边界。
本文的目标是作者要构建一个'精瘦但高性能'的 PyTorch 原生训练框架 Molt，目标是把'从一个关于 agent/奖励/算法的想法'到'一个可信实验'之间的距离压到最短，同时不牺牲大策略所需的性能。具体而言：代码要紧凑干净到研究者能一次读完、AI 编码助手能整体理解并端到端追踪；算法改动应是一次本地编辑而非跨层重构；规模从单机小实验一直到 1T 级 MoE 都用同一条可见的训练回路表达。性能上则要求在显式对齐的协议下与最先进的 Megatron 栈达到统计意义上的可比吞吐——这是'精简'被允许的前提条件，否则简化就不成立。
与已有工作不同的是，本文的独特切入角度是：拒绝'复杂度是强大 RL 基础设施的必然代价'这一默认假设，认为复杂度是从超大规模继承下来的、可以被选择性地放弃。Molt 不重新实现超大规模能力，而是组合已被各自在超大规模上验证过的组件（Ray 做放置、vLLM 做 rollout、NeMo AutoModel + FSDP2 做训练），三者都不 fork，所有上游改进随发版即得。这与'多后端抽象'路线形成鲜明对照——多后端抽象正是层层间接的来源，Molt 选择只支持一个训练后端、一个服务引擎，把这一层彻底删掉而非隐藏起来。其赌注是：整套算法流程可以保持小到能被人类和 AI 助手整体读懂，而评估证明这个赌注不花任何吞吐代价。
核心方法
Molt 的整体直觉是'四个概念、一个回路'：一个 agent（普通 Python，产生动作和奖励）、一个 generator（针对服务引擎做 token 精确采集）、一个 trainer（在单个 FSDP2 policy actor 上的可见训练循环）、以及若干 estimator/loss（关于奖励、group、token trace 的纯函数）。算法上的任何一次改动只触碰这四者之一。技术路线上，运行时就是三个组件加一个异步循环：Ray 提供放置与异步队列，把一个 agent 池、若干 vLLM rollout 引擎和一个基于 NeMo AutoModel 的 FSDP2/EP/CP 策略 actor 连起来；没有混合控制器、没有按后端的适配层、没有独立参数服务器。组件间契约是 token-first 的：token id、每 token 对数概率、动作范围、奖励、多模态张量从引擎采样器一路对齐到 loss，没有任何组件从文本重新派生 token。一个不变量贯穿全回路——Molt 绝不训练它没有生成过的 token。
核心创新是把五条原则压缩成三个可执行的正确性不变量。token identity：被采样的 token id 而非重 tokenize 的转录定义轨迹；policy-version 语义：可训练 token 保留行为策略对数概率 $\log\pi_{\text{beh}}(a_t|s_t)$，异步时做每 token 重要性校正 $\rho_t=\exp(\log\pi_\theta-\log\pi_{\text{beh}})$；forward consistency：rollout 与 actor 必须在模型语义上一致，含多模态扩展与 MoE 路由。与已有方法的本质区别是 token-first 的 agent 边界：Molt 起回环 chat 服务器挡在引擎前，凡跑在标准 OpenAI/Anthropic SDK 上的现成 agent 代码都能原样训练，SDK 流量被解码成逐 token 累积并做 TITO 采集，无需 extra_body、logprobs=true 或会话管道；上下文压缩被自动检测为轨迹分段。这把 Polar、Agent Lightning 重建 token 忠实轨迹的思路内化进了研究 trainer。
方法步骤详情
流程逐步拆解。(1) 作者 agent：用普通 Python 继承 Env 或 ChatAgent 返回 Result(reward=...)；Env 形式框架拥有 LLM 循环（生成/分词/多模态记账/逐轮预算并调 step()），ChatAgent 形式用户拥有循环、框架只起回环服务器并按 ctx.base_url 会话 id 做 TITO 采集。(2) 数据通路：单一数据集服务两种形式，都消费相同 chat-format 数据，由 runner 属性决定模板由谁应用（Env 预渲染、ChatAgent 服务器渲染一次）。(3) 流式池：始终让 prompt group 在飞，凑够即发训练 batch，队列深度解耦训练吞吐与重尾生成延迟。(4) token 精确传输：prompt 以 token id 入、补全以 token id 加 logprob 出，中途不过分词器；路由器把同 rollout 请求钉在前缀缓存所在引擎。(5) 部分 rollout：权重更新不丢在飞请求，暂停→NCCL 广播分片→恢复，混用策略版本时每 token 保留采样 logprob 并套 $\rho_t$ 校正。(6) 权重 refit 经 NCCL 从 actor 直播各引擎。(7) MoE 一致性：引擎返回每 token 专家选择，训练时复现（router replay）。(8) 算法层：估计器按名字选取为纯函数（默认 critic-free REINFORCE++，可切 GRPO/RLOO/Dr.GRPO/GAE+PPO），loss 用全 batch token 均值归一化 $\mathcal{L}_{\text{total}}/N_{\text{unmasked}}$；异步下套序列门 $[0.99,1.01]$ 的每 token 校正。每步报分阶段计时，算法改动同 step 可见。
技术新颖性
新颖性集中在几点。其一，回环服务器把'让现成 SDK agent 训练'从'需集成代码'变成'零集成'，并把上下文压缩当成轨迹分段事件，使长程 agent 仍可训练——把 harness 侧 TITO 能力内化进研究 trainer。其二，'单后端、零 fork'是杠杆而非妥协：不 fork vLLM，引擎升级只是容器版本钉、投机解码/前缀缓存/CUDA graphs 一律按 flag 流入；多后端抽象本就是层层间接的源头，删掉它就删掉了那层，这是 Molt 约 8.6K 行而 verl 约 62K、slime 约 25K 的根因。其三，'规模即配置'：FSDP2 与 AutoModel 原生 TP/EP/CP 组合，DeepSeek-V3 级配置写成 --fsdp.ep_size 256 而非后端迁移，已在 700B MoE、EP256 上端到端跑通，稠密 4B 与超大 MoE 共用同一条 lean 回路。其四，估计器为纯函数、过滤阶段为回路某点的函数加 flag，使'加估计器'收敛为 function+flag+metric+test 四个工件，AI 编码助手据此能独立完成同一编辑。
实验结果
结果围绕三个问题。第一，RL 代码面：import-graph 计数下 Molt 约 8.6K Python 行，对比 verl 约 62K、slime 约 25K、OpenRLHF 约 7.2K——与 OpenRLHF 同档却覆盖到 verl/slime 的超大规模。第二，引擎优化是否以配置到达：Qwen3-30B-A3B（2 节点 × 8 H100，8 训练 + 8 rollout GPU，32K 多轮工具任务）上自动前缀缓存下多轮 re-prefill 仅 0.05s；启用 MTP 投机解码头后每步生成 329s→64s（约 5×），单一配置切换把配方从生成受限变训练受限；--fsdp.offload optimizer 把 actor 峰值显存 64.7GB→46.4GB（省 18.3GB，正 8-GPU 临界差），代价 policy_train 213s→251s（+18%）。第三，与 Megatron 栈正面比较：显式对齐协议（Tab.2）下 Molt（DP8/EP8/TP1）对 slime（TP4+SP/CP1/EP8），均全异步、分到 8+8 GPU、都不加载参考模型，三次独立运行每优化步 119.4±2.3s 对 109.5±10.3s，slime 跨运行带（102–121s）与 Molt 重叠，tok/GPU/s 为 461 对 502，两栈统计可比。训练布局是一阶因素：把 32K 调的 CP 强加到 16K 任务让 Molt 步时增约 30%。需留意公平性说明：该 128 专家 checkpoint 暴露上游分布式-MoE 前向不匹配，actor logprob 与独立参考前向差约 1 nat（$\Delta\approx 1$ nat），$[0.99,1.01]$ 序列门拒绝整批，报告步时仅反映纯吞吐而非有效策略更新，收敛对齐待上游修复。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 代码量（RL 路径，import-graph 计数） | Python LOC | 约 8.6K | verl 约 62K；slime 约 25K；OpenRLHF 约 7.2K | 相对 verl 约缩小 7 倍，与 OpenRLHF 同档却覆盖更大规模 |
| 投机解码加速生成（Qwen3-30B-A3B，每优化步） | 生成墙钟时间 | 64s（MTP head 启用） | 329s（关闭） | 约 5× 加速，单配置切换即从生成受限转训练受限 |
| 优化器 CPU offload 的显存/速度权衡 | actor 峰值显存 / policy_train 时间 | 46.4GB / 251s | 64.7GB / 213s（不 offload） | 省 18.3GB（满足 8-GPU 放得下的临界），代价 +18% 训练时间 |
| 匹配协议下的端到端吞吐（每优化步墙钟） | mean ± s.d.（三次独立运行） | Molt 119.4±2.3s；461 tok/GPU/s | slime（Megatron-Core+SGLang）109.5±10.3s；502 tok/GPU/s | 统计可比，slime 离散带与 Molt 重叠；均值差约 9% 在跨运行噪声内 |
局限与改进
作者坦承的关键局限：35B 工作负载的基准 checkpoint 暴露上游分布式-MoE 前向不匹配——在这个对路由敏感的 128 专家 checkpoint 上 actor 对数概率与独立参考前向相差约 1 nat，于是 $[0.99, 1.01]$ 序列门把整批拒绝，使 Tab.3 报告步时只是'无有效策略更新'的纯吞吐测量；35B 工作负载虽无此问题、门也不滤序列，但收敛性对齐验证仍要等上游修复。其次，作者明确不提供企业数据转换层与统一控制平面（Yan et al. 指出的另两层），多版本 serving、skew 预测器、弹性环境服务、最优资源规划都在范围外，因此 Molt 不是端到端的自我演化部署 agent 系统。从我自己的观察补充：论文缺学习曲线、奖励收敛、最终评测指标（如 pass@n 收敛值）等算法质量结果，正面比较只有 slime 一个对手、只在一个模型/任务上；单后端策略换来了精简却牺牲了部署灵活性（想用 TRT-LLM/SGLang 做 rollout 或 DeepSpeed/Megatron 做训练后端时没有现成路径）。
独立分析的弱点
其一，收敛性证据缺失。实证集中在工程指标（吞吐、显存、代码量、引擎特性），没有一次完整训练曲线、奖励随步数变化或下游评测收敛值，作者自陈等上游 MoE 不匹配修好后才能做收敛对齐——改进方向是补做一组在 35B/真实数学或 agentic 任务上与 slime 完全对齐的训练曲线与最终 pass@n，把'吞吐对齐'升级为'收敛对齐'。其二，泛化面窄：正面比较只有一个对手（slime）、一个模型（Qwen3-30B-A3B）、一个 16K 任务，作者自陈 16K 是最不利于洗出后端差异的设定，改进方向是纳入 verl/NeMo-RL 等更多栈、覆盖稠密模型与更长上下文（32K–128K）以验证'后端差异趋于消失'的预测。其三，单后端的部署刚性：坚持单一训练后端（AutoModel）与单一 rollout 引擎（vLLM）是精简的代价，对想换 SGLang/TRT-LLM 或 DeepSpeed 的用户没有平滑迁移路径，改进方向是提供最小、可选的 backend trait 接口而非整套多后端抽象。其四是生态空白：缺企业数据转换层与控制平面，做生产化自我演化 agent 时需自行补齐。
未来方向
作者明确提出的未来工作有两条主线：一是把 Molt 已在 700B MoE、EP256 上跑通的'全异步回路（rollout+refit+optimizer step）'推进到端到端收敛测量，目标迈向并越过 NVIDIA GB300 上的 3 万亿参数标记，并与 NeMo AutoModel 团队继续推进可训练的模型尺寸；二是把 Molt 作为高可靠性的研究基底，开展质量与可用性（quality/usability）的用户研究，验证'人类可读性优先 + AI 编码助手可追踪'的设计是否真的加速研究迭代。基于本文成果可延伸的方向还有：把 Molt 的 token-first 契约与外部 harness 的轨迹协议（Yan et al. 提出的 step-granular trajectory protocol）对接，作为互补的轨迹来源；引入 skew 预测器、多版本 serving、弹性环境调度以覆盖作者主动让出的调度类贡献；把回环服务器的 TITO 采集与上下文压缩分段能力推广到 OSWorld、浏览器自动化等更长程 agent，研究分段边界对信用分配的影响。
复现评估
复现性整体良好。Molt 以 Apache-2.0 开源，仓库提供完整框架、每项测量背后的参考 agent 与 one-command recipes 及预构建容器。Tab.2 钉死被比较栈 commit：Molt cb2cae11、slime 5d7296a7、Megatron 1dcf0daf；附录 A.2 给出 Qwen3.6-35B-A3B geo3k 配方（2 节点 × 8 H100，8 训练 + 8 rollout GPU，32K 上下文，CP8/EP8/TP1，4 prompts × 4 samples，temperature 1.0，seq-mask-tis，R3 on）。scaling knobs（Tab.4）与 agent 代码示例完整，照 CLI 即可起跑。主要门槛是算力：正面比较至少需 2 节点 × 16 张 H100、700B MoE 需 EP256 即大量 GPU，普通研究者难直接复现大规模实验；小规模（4B 稠密、单节点）相对可行。作者诚实标注 35B checkpoint 的上游 MoE 不匹配会拒绝整批、使报告数字只反映吞吐而非有效更新，降低了训练有效性可复现度，需等上游修复。
论文图表