One-sentence Summary

NVIDIA researchers introduce Molt, a compact PyTorch-native training framework for agentic reinforcement learning that uses an asynchronous loop to train multimodal and mixture-of-experts policies while ensuring models never train on tokens they did not generate, achieving statistical comparability to state-of-the-art Megatron-based systems and providing a codebase compact enough for a human to hold in their head and for an AI coding assistant to read and reason about in its entirety.

Key Contributions

• Molt provides a PyTorch-native training framework for agentic reinforcement learning whose compact codebase is designed to be read and reasoned about in its entirety by a researcher or an AI coding assistant.
• The framework enforces token-level consistency by training multimodal and mixture-of-experts policies only on self-generated tokens, preserving behavior log-probabilities, and using router replay to prevent MoE instability from divergent rollout and training routes.
• Under a matched fully asynchronous protocol, Molt achieves statistical parity with a state-of-the-art Megatron-based stack and scales the same lean loop from 4B to 700B MoE parameters at expert parallelism 256, with an open-source release that includes recipes, containers, and one-command reproduction of results.

Introduction

The authors address a growing tension in agentic reinforcement learning research: mainstream RL frameworks are architected for hyperscale training workloads and impose layers of distributed backend, controller, and configuration complexity that slow down the rapid algorithm modification central to research. Prototyping a new idea often requires tracing changes through many glue layers, and subtle numerical mismatches between serving and training, such as tokenization or expert routing divergences, can silently bias gradients without raising errors.

To reduce this friction, the authors introduce Molt, a lean, PyTorch-native framework that prioritizes human and AI coding assistant readability as a primary design constraint. Its main contribution is a framework with a codebase several times smaller than comparable Megatron-based production stacks that matches their throughput while enforcing three core correctness invariants: training on exactly the generated token IDs, preserving behavior-policy log-probabilities, and ensuring rollout-actor forward consistency.

Method

The authors construct Molt as a streamlined system for agentic reinforcement learning that deliberately eschews the layered indirection typical of hyperscale training stacks. The entire runtime is organized around a single asynchronous loop that connects three well-defined component groups: an agent pool, a set of vLLM rollout engines, and one trainable policy actor built on NVIDIA’s AutoModel with Fully Sharded Data Parallel 2 (FSDP2), expert parallelism (EP), and context parallelism (CP). A Ray-based queue decouples generation from training, keeping prompt groups in flight and emitting a training batch as soon as enough completions arrive. This architecture avoids any kind of adapter layer or parameter server, and the invariant that the system never trains on a token it did not generate is enforced by a token-first contract: every message carries token ids, per-token log-probabilities, and aligned multimodal tensors, and no component ever re-tokenizes text.

The agent boundary respects a principle of modularity along the algorithm’s own objects. A researcher provides an AgentRunner module, and the rest is ordinary Python. In Env mode, the framework drives the LLM loop, calling the user’s step function after each model action and handling tokenization and budgets. In ChatAgent mode, existing agent code that speaks OpenAI or Anthropic wire protocols is served verbatim: Molt launches a loopback chat server that decodes every request into a token-exact accumulation, delivering token-in/token-out capture without any integration code. Context compaction (summarizing or dropping earlier turns) is automatically handled by segmenting trajectories, preserving end-of-rollout rewards. Both agent forms consume the same chat-template data, and a single dataset serves either path without data-path confounding.

Transport is kept correct by construction. The system carries no vLLM forks; all interactions exploit stable, documented endpoints. Trainer-facing generation uses the engines’ token-level interface, with prompts and completions as token ids along with per-token log-probabilities. The request router pins all requests of one rollout to the same engine to preserve prefix caching, and multimodal prompts are rendered server-side to avoid silent image loss. Weight updates are handled by an NCCL broadcast directly from the actor to every engine, bypassing the router entirely.

Scale is expressed through configuration rather than migration. The same launch script that trains a dense 4B model also expresses a large Mixture-of-Experts (MoE) configuration via straightforward parallelism flags (e.g., --fsdp.ep_size 256). For MoE architectures, rollout routing replay is employed locally: the engine returns its per-token expert choices, and the actor replays them during training, eliminating the failure mode where rollout and training make different routing decisions. Other optimizations such as speculative decoding, prefix caching, and CUDA graphs surface as engine flags, while optimizer offloading to host memory and conservative guard rails reject incompatible parallelism combinations at startup.

The algorithm layer similarly avoids inheritance hierarchies. Advantage estimators are pure functions of rewards and groups, selected by name (REINFORCE++ with group-mean baseline, RLOO, GRPO, GAE with a critic, or on-policy distillation). One global loss normalization uses the whole-batch token mean, making the update invariant to data-parallel size and gradient accumulation. An off-policy correction with a sequence-level gate, together with dynamic batch filtering, handles asynchronous rollouts and ensures that every token’s training log-probability matches the generation-time value. The end result is a system where an algorithmic change, such as introducing a new advantage estimator, touches only a single function, a flag, and a logged metric, and the entire loop remains readable in one pass.

Experiment

The evaluation quantifies the cost of leanness by measuring the RL stack’s footprint, the throughput impact of engine optimizations, and its performance against a Megatron-based baseline. The framework spans only 8.6K lines, and upstream engine features (prefix caching, speculative decoding, optimizer offloading) arrive as configuration flags, delivering clear speedups and memory savings. A controlled head-to-head comparison on a 30B MoE model shows statistically comparable per-step times (119.4 ± 2.3 vs. 109.5 ± 10.3 s), confirming that the composed, lean design achieves throughput parity with a full stack without adding performance overhead.

Under a matched asynchronous protocol on disaggregated GPUs, Molt and slime achieve statistically comparable per-optimizer-step times, with overlapping run-to-run variability and no clear winner. Training layout is a first-order factor: mismatched parallelism settings can inflate step time by roughly 30% on this 16K workload. A benchmark-side MoE forward mismatch prevents effective policy updates for the 128-expert checkpoint, so reported step times measure throughput only; a 35B workload shows no such gap. Molt and slime step times overlap within run-to-run variability, indicating no significant backend performance difference under the controlled protocol. Forcing a context-parallel layout optimized for 32K contexts onto a 16K workload slows Molt’s step time by around 30%. A distributed-MoE log-probability mismatch in the benchmark checkpoint causes the sequence gate to reject the batch, so optimizer steps measure throughput without effective policy updates. As output lengths grow toward reasoning or agent regimes, generation dominates step time and the training backend’s relative contribution shrinks toward irrelevance.

A head-to-head comparison between the lean Molt stack and the Megatron-based slime stack on Qwen3-30B-A3B with 16K context shows statistically comparable optimizer step times (119.4 ± 2.3 s versus 109.5 ± 10.3 s). The roughly 9% mean difference falls within slime's cross-run variability, indicating no throughput superiority for either. At longer output lengths the training backend's share of total step time shrinks, reinforcing the practical equivalence. Molt (FSDP2 + vLLM) and slime (Megatron-Core TP4+SP + SGLang) deliver comparable per-step throughput when each uses its recommended parallel layout on 8+8 H100 GPUs. slime's step time spread of 102–121 seconds fully envelopes Molt's 117–122 second band, making the difference statistically insignificant. Both stacks run fully asynchronously with disaggregated rollout and training, and neither loads a reference model, ensuring identical per-step algorithmic work. A distributed-MoE forward mismatch in the benchmark checkpoint causes actor log-probabilities to deviate by about one nat, so reported step times reflect throughput alone without an effective policy update.

Under a matched protocol with disaggregated rollout and training, the two stacks achieve throughput parity: step times are 119.4 ± 2.3 s for MOLT and 109.5 ± 10.3 s for slime, with a mean difference of roughly 9% that lies within cross-run variability. Both stacks overlap generation with training, so end-to-end step time is the meaningful throughput metric, and any residual backend advantage shrinks as generation length increases. The reported tokens-per-GPU-second values reflect only generation throughput because a checkpoint mismatch prevented effective policy updates. Step time distributions overlap across runs, indicating no statistically meaningful speed advantage for either stack. Slime yields 502 tokens/GPU/s versus 461 for MOLT, but the spread in slime run times (102–121 s) encompasses the MOLT mean. Generation and training are fully overlapped on disaggregated GPUs, making step time the end-to-end throughput metric. The 9% mean step-time difference falls within cross-run variability, so neither stack claims superiority. As output lengths grow toward reasoning-scale trajectories, backend-specific overhead becomes negligible relative to generation time. The 16K context length is the regime least favorable to hiding backend differences, yet the stacks are already comparable. Forcing a context-parallel layout tuned for 32K contexts inflates MOLT's step time by roughly 30% at this sequence length. A distributed-MoE checkpoint mismatch caused actor log-probabilities to diverge by about one nat, making the measurements throughput-only without an effective policy update.

A head-to-head comparison of the Molt (FSDP2 + vLLM) and slime (Megatron-Core with SGLang) stacks under a matched asynchronous protocol on disaggregated GPUs shows statistically indistinguishable step times, with overlapping run-to-run variability and no backend achieving a meaningful speed advantage. Training parallelism layout acts as a first-order factor: applying a context-parallel configuration tuned for 32K contexts to a 16K workload inflates step time by roughly 30%. A distributed-MoE forward mismatch in the 128-expert benchmark prevents effective policy updates, so the reported step times measure throughput only, while a 35B workload exhibits no such gap. As output lengths grow toward reasoning or agent regimes, generation dominates total step time and the training backend’s relative contribution shrinks to practical irrelevance, reaffirming the equivalence of both stacks.