Molt: Training Agents the Right Way

The Paper

"Molt: A Scalable PyTorch-Native Training Framework for Agentic Reinforcement Learning" was published in July 2026 by Jian Hu, Huiying Li, Hao Zhang, Binfeng Xu, Yifan Zhang, Shaokun Zhang, Hemil Desai, Michael Demoret, Pavlo Molchanov, Jan Kautz, and Yi Dong from NVIDIA. The central claim is that a fully asynchronous agentic RL training loop can be expressed as an ordinary Python program - compact enough for a researcher to read end-to-end, PyTorch-native with no heavy distributed backend dependencies, and statistically comparable in throughput to state-of-the-art Megatron-based stacks.

Read the Paper on arXiv →

The Problem Before This Paper

Existing agentic RL frameworks - verl (~62K RL LOC, FSDP2/Megatron), OpenRLHF (~7.2K LOC, DeepSpeed ZeRO-3), and slime (~25K LOC, Megatron) - require researchers to route every algorithm change through multiple abstraction layers: trainers, distributed backends, rollout glue code, and parameter servers. Adding a new estimator or modifying the policy gradient target means touching subsystems that were designed for production breadth, not research iteration. Beyond ergonomics, the deeper problem is correctness: most frameworks allow subtle policy-version drift in asynchronous settings, where tokens generated under one checkpoint can end up in a training batch attributed to a different policy. This mixes behavior and training signal in ways that are hard to detect and can produce unstable or misleading gradient estimates, particularly in multi-turn agentic tasks where the policy interacts with tools or generates long contexts.

What They Built

Molt is structured around four primitives: Agent (plain Python, produces actions and rewards), Generator (token-exact capture against vLLM serving engines), Trainer (single FSDP2 policy actor), and Estimators (pure functions of rewards, groups, and the token trace). The asynchronous loop maintains a streaming pool of prompt groups - all rollout samples of one prompt - in flight continuously, emitting a training batch when enough groups complete. Weight updates do not discard in-flight requests: the system pauses engines, broadcasts actor shards via NCCL directly to each engine bypassing the request router, then resumes retained requests. Every action token retains its log-probability from sampling time; per-token importance correction is applied at loss computation, and partial rollout is refused without this correction enabled. Molt exposes two agent interfaces: an Env form where the framework drives the full LLM loop, and a ChatAgent form where the researcher's own code uses a stock OpenAI or Anthropic SDK against a loopback chat server - every SDK request is decoded server-side into token-exact accumulation via token-in/token-out (TITO) capture, requiring no extra_body or session plumbing. Both interfaces share one data path. The algorithm layer is deliberate in its flatness: estimators are selected by flag (--algo.advantage.estimator), implemented as pure functions with no strategy classes or inheritance. Supported estimators include REINFORCE++, GRPO, Dr. GRPO, RLOO, GAE+PPO critic, and on-policy distillation. Loss normalization uses a global whole-batch token mean - a single denominator shared by policy-gradient, KL, and entropy terms - making updates invariant to data-parallel size and gradient-accumulation depth.

Key Findings

• Three correctness invariants enforced by construction. Token identity (sampled token IDs define the trajectory, not retokenized transcripts), policy-version semantics (asynchronous tokens carry their behavior-policy log-probabilities and are explicitly corrected at training), and forward consistency (rollout and actor must agree on model semantics including MoE routing and multimodal expansion). Partial rollout without importance correction is refused, not warned.
• MoE routing replay prevents silent training divergence. For mixture-of-experts models (DeepSeek-V3-class, Qwen3-30B-A3B), the rollout engine returns per-token expert choices; the actor replays them during training via vLLM's native route capture and AutoModel RouterReplay. Without this, the sparse computation graph diverges between rollout and training, producing corrupted gradient estimates that are invisible in aggregate loss metrics.
• Speculative decoding reduces per-step generation time by ~5x. On Qwen3.6-35B-A3B (2 nodes, 32K multi-turn tool-use), enabling the MTP head dropped per-step generation from 329 seconds to 64 seconds, shifting the workload from generation-bound to training-bound. This is the largest single-knob throughput lever in the framework.
• Optimizer CPU offload saves 18.3 GB peak GPU memory at +18% training time cost. On the same 35B recipe with an 8-GPU training partition, --fsdp.offload_optimizer reduces peak from 64.7 GB to 46.4 GB. The tradeoff fits multi-node configurations where VRAM is the binding constraint but training time is not.
• Scale verified at 700B MoE with expert parallelism 256. The same async loop runs on a 700B mixture-of-experts model and a 4B dense model via configuration change only. No code path splits between scale tiers.

Results

Head-to-head against slime (Megatron-Core + SGLang) on Qwen3-30B-A3B, 2x8 H100s, 8+8 disaggregated rollout, fully asynchronous protocol: Molt achieved 119.4 +/- 2.3 seconds per step at 461 tokens/GPU/second, versus slime at 109.5 +/- 10.3 seconds at 502 tokens/GPU/second. The mean difference is approximately 9%, within slime's cross-run variance range of 102-121 seconds. The authors note that forcing context-parallel degree tuned for 32K onto this 16K workload inflates Molt step time by roughly 30%, and that training layouts differ (Molt: DP8/FSDP2, EP8, TP1; slime: TP4+SP, CP1, EP8). A disclosed benchmark integrity issue - actor log-probabilities on the 30B checkpoint diverged from reference forward by approximately 1 nat due to a distributed-MoE forward mismatch, causing the [0.99, 1.01] sequence gate to reject batches - means these numbers measure raw throughput only, not effective policy gradient updates on that checkpoint. At ~8.6K RL LOC against verl's ~62K, the codebase differential is roughly 7x. Prefix caching on cache hit re-prefills in 0.05 seconds.

Why This Matters for AI and Automation

• The correctness gap is silent. Policy-version drift and MoE routing divergence do not produce obvious training failures - they produce subtly wrong gradient estimates that accumulate across steps. In multi-turn agentic tasks with tool calls, this compounds: each turn's gradient depends on log-probabilities that were never accurately recorded in frameworks without TITO capture.
• ChatAgent form makes existing agent code trainable without modification. If you are building agents with the Anthropic or OpenAI SDK today - the same interfaces used in production agentic deployments - Molt's loopback chat server can capture those trajectories token-exactly for GRPO or REINFORCE++ training. There is no requirement to rewrite agent logic to fit a framework-specific interface.
• Comprehensibility as a design constraint is underrated. Molt's ~8.6K LOC for the full RL stack means a practitioner can trace a gradient estimate from reward to optimizer step without crossing abstraction boundaries they do not own. For research teams iterating on estimators or loss formulations, this reduces the gap between "idea" and "running experiment" from days to hours.
• This connects directly to Week 24's invisible reasoning finding. Baherwani et al. showed that models can perform consequential computation in tokens that carry no semantic trace in the output. Molt's token identity invariant - training only on the exact token IDs the policy sampled - matters more in this context: a framework that retokenizes transcripts or drifts policy versions trains on a corrupted signal that cannot be distinguished from correct training without explicit auditing. Molt makes this class of error structurally impossible rather than merely warned against.

My Take

The benchmark comparison with slime is interesting but not the paper's actual contribution. A 9% mean difference in throughput with overlapping variance ranges, on a workload where the authors disclose the run encountered a correctness failure, is not a performance claim worth anchoring on. What Molt actually demonstrates is that the correctness invariants the field has been treating as implementation details - token identity, policy-version semantics, forward consistency - can be made structural guarantees in a framework small enough to read. The ChatAgent form is the underemphasized result: the ability to train agents written against stock SDK interfaces without any framework coupling is a meaningful unlock for teams building production agentic systems. The disclosed benchmark issue is also worth noting: the fact that a 1 nat divergence in actor log-probabilities on a distributed-MoE forward triggered the importance correction gate rather than silently producing wrong gradients is itself a demonstration of the framework working as designed. Most frameworks would have continued training. The open question is whether the comprehensibility argument holds as the estimator surface grows - Molt currently supports seven estimators, which is not a large surface. What happens to the LOC count and the "ordinary program" framing when the field standardizes on twelve estimators with conditional interactions? That is not a criticism of this paper, it is the natural lifecycle question for any framework that prioritizes readability over breadth.

Discussion question: Molt's ChatAgent form allows agents built against stock Anthropic or OpenAI SDK interfaces to be trained with GRPO or REINFORCE++ without code changes. If you are building production agentic systems today using those APIs, what would you need to change about your reward signal design to make that training loop produce useful gradient updates - and is the primary bottleneck the framework, the reward signal, or the evaluation harness?