Source-linked AI summary
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
TL;DR
Agentic RL research needs infrastructure that makes frequent algorithm changes easier to trace and implement without sacrificing the performance demanded by large policies. Molt addresses this with a compact, PyTorch-native asynchronous framework whose performance is statistically comparable to a state-of-the-art Megatron-based stack under a matched protocol.
Problem
Mainstream agentic RL frameworks make routine algorithm changes costly by spreading them across trainer, distributed-backend, and rollout-engine layers.
Method
Molt uses a compact, readability-first PyTorch-native design with one asynchronous loop, ordinary Python agents, token-first boundaries, and minimal un forked backend abstractions.
Results
Molt is statistically comparable in system performance to a state-of-the-art Megatron-based stack under a matched protocol.
Takeaways & Limitations
Molt supports research workflows in which environments remain ordinary programs and the complete RL path stays small enough for humans and AI coding assistants to inspect.
Takeaways & Limitations
On the benchmark’s routing-sensitive 128-expert checkpoint, an upstream distributed-MoE forward mismatch makes the reported rows throughput-only.
Abstract
from arXiv · showhide
Agentic reinforcement learning research is constant algorithm modification, new estimators, new pipeline stages, new rollout schemes, and in mainstream frameworks each change threads through layers of trainer, distributed backend, and rollout glue: the cost lands on the researcher at every iteration. Molt is a PyTorch-native training framework built to keep that cost small: a codebase compact and clean enough for a researcher to hold in their head, and for an AI coding assistant to read and reason about in its entirety, so the algorithm flow can be traced and changed end to end. The agent is an ordinary program, and one asynchronous loop trains multimodal and mixture-of-experts policies while never training on a token it did not generate, consistent in tokens, policy versions, and model semantics. Leanness does not cost performance: under a matched, fully asynchronous protocol, Molt is statistically comparable to a state-of-the-art Megatron-based stack. Molt is open source and provides recipes and containers at https://github.com/NVIDIA-NeMo/labs-molt.
1. Introduction
Molt argues that agentic RL infrastructure should minimize the distance from algorithmic ideas to trustworthy experiments without sacrificing large-policy performance. It pursues this through a compact PyTorch-native framework whose readable, modular design preserves correctness across asynchronous, multimodal, and MoE training.
- Motivation: Mainstream RL frameworks make routine algorithm changes costly by spreading them across trainer, distributed-backend, and rollout-engine layers.Agentic RL research repeatedly modifies estimators, experience filtering, and rollout schemes.
- Motivation: Molt makes source code part of the research interface, enabling researchers and AI coding assistants to trace samples from agent invocation through policy loss.The framework targets a short path from an idea about an agent, reward, or algorithm to a trustworthy experiment.
- Framework: Molt combines Ray, vLLM, and NeMo AutoModel in one disaggregated asynchronous loop without forking them, while supporting ordinary Python environments and existing agent SDKs.It can drive an Env or let an existing agent use a stock OpenAI or Anthropic SDK through a loopback capability.
- Correctness: Three correctness invariants govern Molt: token identity uses sampled token ids, policy-version semantics preserve behavior-policy log-probabilities with asynchronous correction, and forward consistency aligns rollout and actor model semantics.Forward consistency includes multimodal expansion and mixture-of-experts routing.
- Thesis: Molt’s central thesis is that capable RL infrastructure need not be complex: its algorithm flow remains small enough for humans and AI assistants to read as a whole, with no throughput cost in evaluation.The framework’s distinct combination is readability, compactness, correctness, modularity, and performance parity rather than one isolated mechanism.
2. Design Principles
Molt’s design addresses the mismatch between hyperscale infrastructure and rapid algorithmic iteration by favoring readable composition over hyperscale-specific layering. Five principles govern the framework: human- and AI-readable code, minimal backend structure, state-of-the-art performance, algorithm-aligned modularity, and detailed correctness guarantees.
- Design rationale: Molt composes separately frontier-hardened components so scalability is inherited without re-implementing hyperscale-specific layering.This approach targets research iteration while covering the same scalability band as hyperscale stacks.
- Principle 1: Readability: Every function must expose understandable control and data flow, enabling researchers and AI coding assistants to trace features end to end.Traceability extends from CLI flags through executed branches, tensors, metrics, and tests; unnecessary indirection is treated as a defect.
- Principle 2: Minimal code: Molt minimizes code and deliberately uses exactly one training backend, AutoModel, and one serving engine, vLLM, without forks.Deletion is preferred over addition, while performance controls and observability are retained and upstream improvements require no rebasing.
- Principle 3: Performance: Performance parity with a state-of-the-art Megatron-based stack under a matched protocol is a design constraint, not an optional outcome.Leanness cannot shift time onto the training path, and composed components must preserve the engines’ own optimizations.
- Principle 4: Modularity: Molt’s components map one-to-one onto RL algorithm objects rather than infrastructure layers, so algorithmic edits touch exactly one component.The principle rejects adapter layers and plugin registries, with the agent contract and pure-function estimators serving as examples.
- Principle 5: Correctness: Molt treats numerical fidelity as a first-class guarantee by training only on generated tokens, aligning log-probabilities, and gating sequence-level mismatch.Residual training–inference mismatch is monitored at every step behind a hard sequence-level gate.
3. The System: Four Concepts, One Loop
Molt reduces agentic RL system complexity to four code-mapped concepts implemented through one visible, asynchronous loop. Its modular agent interfaces, token-exact transport, and configurable parallelism preserve algorithmic flexibility while supporting multimodal and MoE training.
- Four Concepts, One Loop: Four load-bearing concepts map one-to-one onto code: agent, generator, trainer, and pure-function estimators and losses.An algorithmic change touches exactly one of the four concepts.
- Four Concepts, One Loop: The runtime uses three components and one loop: an agent pool, vLLM rollout engines, and one FSDP2 policy actor connected by an asynchronous queue.Ray provides placement and queueing; reference workers and a PPO critic are optional.
- Agent Interfaces: ChatAgent supports stock OpenAI or Anthropic SDK-based agents without integration code through loopback-server token-in/token-out capture.The server accumulates one token-exact trajectory and eliminates retokenization drift.
- Transport: Molt carries no vLLM patches, making engine upgrades container pins while relocating token-exactness requirements to stable client-side interfaces.Trainer-facing generation uses token ids and per-token log-probabilities, avoiding tokenizer-mediated text during episodes.
- Production Features: 1T-class MoE is expressed by configuration: a DeepSeek-V3-class setup uses –fsdp.ep_size 256 rather than a backend migration.FSDP2 composes with AutoModel-native tensor, expert, and context parallelism and matching vLLM controls.
- Production Features: Rollout routing replay makes MoE training consistent by returning per-token expert choices for the actor to replay.This locally prevents rollout and training routers from evaluating different sparse computation graphs.
4. Evaluation: Does Leanness Cost Throughput?
Molt tests whether a compact composed design preserves practical optimizations and throughput. On a matched asynchronous protocol, it is statistically comparable to a Megatron-based stack while retaining substantial serving and memory benefits.
- Evaluation scope: The evaluation measures RL surface area, optimization propagation, and throughput parity against a state-of-the-art Megatron-based stack.The protocol compares Molt with slime under shared pinned settings, backend-recommended layouts, and three independent runs.
- Serving and memory optimizations: 0.05 s re-prefill demonstrates automatic prefix caching and session-consistent routing on a cache hit.The passage explicitly notes that no cache-miss baseline is available, so this is not an end-to-end speedup claim.
- Serving and memory optimizations: 5× generation reduced per-step generation time from 329 s to 64 s with the checkpoint’s MTP head enabled.A single configuration change moved the recipe from generation-bound to training-bound.
- Serving and memory optimizations: 18.3 GB for 18% lower actor peak GPU memory reduced usage from 64.7 GB to 46.4 GB, while policy_train time rose from 213 s to 251 s (+18%).The offload configuration made the 8-GPU training partition fit.
- Throughput parity: 119.4±2.3 versus 109.5±10.3 s per optimizer step shows statistical comparability between Molt and slime over three runs.The slime cross-run spread was 102–121 s and overlapped Molt’s band; both stacks overlapped generation with training on disaggregated GPUs.
- Throughput parity: At longer 32K–128K reasoning and agent outputs, generation should dominate step time while the training backend’s share shrinks.This is presented as an expected residual-difference trend; the controlled 16K setting is described as least favorable to that effect.
5. Related Work
Molt is positioned against RL training frameworks, asynchronous rollout systems, and agent-trajectory capture systems. Its distinguishing combination is a readability-first research surface, token-first agent contract, and high-performance PyTorch-native implementation, without proposing a new RL objective.
- RL training frameworks: Molt contrasts with hybrid, distributed, Megatron-based, and lightweight PyTorch RL training frameworks by prioritizing a compact, readable research surface.The comparison includes HybridFlow/verl, OpenRLHF, NeMo-RL, TRL, and Slime.
- Asynchronous and disaggregated RL systems: Prior asynchronous and disaggregated systems address staleness, pipeline or skewness bubbles, weight distribution, policy versions, fixed-weight prompt groups, and heterogeneous trajectory scheduling.Examples include AReaL, StreamRL, Laminar, DORA, RolloutPipe, Relax, and RollArt.
- Training–inference consistency: Molt treats training–inference consistency as requiring replayed MoE expert routes, token-level behavior probabilities, and fail-fast checks for unsupported feature combinations.It implements route replay through vLLM native route capture and AutoModel RouterReplay.
- Agent trajectories and harness integration: Agent Lightning and Polar decouple agent execution or harnesses from RL training through data interfaces, hierarchical credit assignment, traffic proxies, and token-faithful trajectory reconstruction.Polar supersedes ProRL Agent and treats an arbitrary harness as a black box.
- RL objectives: Molt proposes no new objective, instead keeping aligned tokens, behavior log-probabilities, complete prompt groups, and consistent sparse routing explicit for established estimators.Its estimators follow REINFORCE++, GRPO, DAPO, and GSPO.
6. Future Work
Future work extends Molt toward end-to-end convergence measurements at and beyond the 3-trillion-parameter scale on NVIDIA GB300 systems. The framework already runs its complete asynchronous training loop on a 700B mixture-of-experts model at expert parallelism 256.
- Future Work: Molt runs the full asynchronous loop end to end on a 700B MoE at expert parallelism 256, using the same lean loop as at 4B.The loop includes rollout, weight refit, and optimizer steps.
- Future Work: Toward and beyond the 3 trillion-parameter mark on NVIDIA GB300, the remaining work is to measure end-to-end convergence.Molt’s configuration surface expresses DeepSeek-V3-class models, while vLLM and AutoModel provide large-scale MoE support and recipes.
7. Conclusion
Molt is designed as a compact, inspectable PyTorch-native substrate for agentic reinforcement learning research, with an explicit end-to-end path that AI coding assistants can trace. Its measurements support comparable system performance under a matched protocol, while its open-source release includes the complete framework, recipes, reference agents, and prebuilt containers.
- Research-oriented design: Molt keeps environments ordinary programs and the complete reinforcement-learning path small enough for researchers and AI coding assistants to inspect end to end.The explicit path spans request tokens, rollout, reward, weight version, and loss, while stable upstream components provide serving and distributed execution.
- Performance and usability: Statistically comparable system performance to a state-of-the-art Megatron-based stack is achieved under a matched protocol.Engine optimizations and optimizer offload arrive as configuration through the composed path.
- Performance and usability: Readable infrastructure makes failures visible, localizable, and difficult to ignore when something anywhere in the stack is wrong.Quality and usability studies are identified as natural next steps for Molt’s high-performance substrate.
- AI-assisted research: Molt is infrastructure designed for research with AI coding assistants, using a codebase sized to be read whole, one visible loop, and parts shaped like the algorithm.The framework presents these properties as a form more research infrastructure may adopt.
- Availability: Molt is open source under the Apache-2.0 license, with the complete framework, reference agents, one-command measurement recipes, and prebuilt training and serving containers.The repository ships the materials behind every reported measurement.
A. Appendix · A.1. Scaling Knobs
Molt’s scaling design uses a shared launch path across dense and mixture-of-experts models. Its principal controls compose actor-side parallelism with FSDP2 sharding, while vLLM mirrors the configuration per engine.
- A.1. Scaling Knobs: A shared launch path spans a dense 4B model and a configured 1T-class MoE.This indicates one launch mechanism covers both model configurations.
- A.1. Scaling Knobs: The same launch path supports both dense and MoE model regimes.The passage names a dense 4B model and a configured 1T-class MoE as examples.
- A.1. Scaling Knobs: Molt provides principal scaling controls for its training setup.These controls are summarized in Table 4.
- A.1. Scaling Knobs: Actor-side parallelism composes with FSDP2 sharding.The two mechanisms form the actor-side scaling configuration.
- A.1. Scaling Knobs: The vLLM side mirrors actor-side parallelism per engine.The mirroring applies to each vLLM engine.
- A.1. Scaling Knobs: The scaling controls are presented as Molt’s principal configuration knobs.Table 4 specifically labels them as scaling knobs.
A.2. Reproduction Note
The reproduction measurements use the shipped Qwen3.6-35B-A3B geo3k RL recipe under a specified distributed setup and training configuration. Head-to-head benchmarks follow Tab. 2, with pinned commits and single-command launches including containers.
- A.2. Reproduction Note: Measurements use the shipped Qwen3.6-35B-A3B geo3k RL recipe on 2 nodes with 8 H100 GPUs each, split between training and rollout.The configuration allocates 8 training and 8 rollout GPUs.
- A.2. Reproduction Note: The recipe uses 32K context, CP8 / EP8 / TP1, 4 prompts × 4 samples, temperature 1.0, seq-mask-tis, and R3 enabled.These settings define the reported single-framework measurement configuration.
- A.2. Reproduction Note: The head-to-head benchmark follows Tab. 2 with framework commits pinned, and each configuration includes a single-command launch recipe and container specifications.The launch recipes and containers are provided in the repository.