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

Molt is a lean, PyTorch-native RL framework that matches hyperscale throughput while keeping the entire algorithm flow readable.

How can we build a reinforcement learning framework for agents that avoids the complexity of custom engine forks and proprietary rollout pipelines?

Mainstream reinforcement learning frameworks are built for hyperscale production, forcing researchers to navigate complex, multi-layered abstractions just to test simple algorithmic changes. Molt replaces these layers with a single, asynchronous PyTorch loop that treats the agent as ordinary Python code and keeps the entire training path small enough to read in one pass. On a 35B multimodal mixture-of-experts workload, Molt achieves throughput statistically comparable to state-of-the-art Megatron-based stacks while maintaining a codebase several times smaller.

Paper Primer

Molt achieves this by composing hardened, upstream components—specifically NVIDIA AutoModel and vLLM—without forking them. The framework hinges on a "token-first" contract: it captures exact token IDs and log-probabilities directly from the serving engine, ensuring that the trainer and rollout engine remain perfectly aligned without needing reconciliation layers.

Molt delivers throughput parity with Megatron-based production stacks.

Head-to-head benchmark on a 35B multimodal MoE model (16K context) showed 119.4 ± 2.3s per optimizer step for Molt versus 109.5 ± 10.3s for the Megatron-based slime stack. The ~9% mean difference falls within the range of cross-run variability, demonstrating no significant performance penalty for the lean design.

The framework significantly reduces the cognitive and maintenance overhead for researchers.

The entire RL path is approximately 8.6K lines of code, compared to 62K for verl and 25K for slime. This reduction allows a researcher to trace a single sample from agent invocation to policy loss in one pass, making the code itself the primary interface for experimentation.

Why does this framework prioritize readability over the multi-backend abstractions found in other stacks?

The authors argue that for research, the cost of understanding and modifying the framework often exceeds the cost of expressing the hypothesis itself. By refusing multi-backend abstraction, Molt eliminates the "layered indirection" that makes hyperscale stacks difficult to hack.

Does the "lean" design limit the ability to scale to large models?

No; scalability is inherited from the underlying components. Molt uses the same configuration surface to express models ranging from a 4B dense model to a 700B mixture-of-experts model, with parallelism knobs (like expert and context parallelism) mapped directly to the training loop.

Molt is designed for an era where AI coding assistants (e.g., Claude Code) share the research workload; by keeping the control flow explicit and the codebase small, the framework allows these assistants to reason about the entire training path without reconstructing hidden registries.

Researchers can now iterate on agentic RL algorithms without inheriting the complexity of hyperscale infrastructure, using a framework that treats code readability as a first-class performance constraint.

The Bottleneck in Agentic RL

We expose how monolithic RL stacks hinder rapid research and motivate a modular PyTorch-native alternative.

Agentic reinforcement‑learning research constantly rewrites estimators, pipeline stages, and rollout schemes; in mainstream stacks each tweak must be propagated through layers of trainer code, distributed back‑ends, and rollout glue, turning a simple idea into a costly engineering effort.

Current agentic‑RL pipelines are built on monolithic, rigid frameworks that demand custom engine forks and tangled integration for every algorithmic change, inflating the effort required to prototype and evaluate new ideas.

The key shift is from monolithic, hyperscale‑oriented frameworks to a compact, PyTorch‑native modular design that preserves performance while dramatically lowering research overhead.

Design Principles for Scalable RL

We outline Molt’s five design principles that enable readable, modular, high‑performance agentic RL.

The bottleneck identified earlier is the heavyweight, monolithic RL scaffolding that forces every new algorithm to fork a custom engine. Molt attacks this pain by insisting that the framework itself be as readable as a tutorial and as modular as a Lego set, without sacrificing performance.

Think of the codebase as a well‑labeled toolbox: both a human developer and an AI coding assistant can locate the exact wrench they need without consulting a separate diagram.

Parse flag → value = “dense”.

Conditional branch checks if value == “dense”. It passes, so the dense backend is instantiated.

No else‑branch is executed because the “moe” path does not exist.

Training loop runs using the single backend’s optimizer and scheduler.

With only one backend the code path is linear, eliminating the hidden branching that previously required a second pass to understand.

How does Molt’s modularity differ from the plugin registries used in other RL frameworks?

Plugin registries expose a generic “register‑component” hook that any file can attach to, creating a sprawling graph of indirect calls. Molt’s modularity is one‑to‑one: each RL object (agent, rollout, loss) maps directly to a concrete Python module, so the call graph is explicit and traceable without a registry.

**Table 1.** Framework comparison: the lean point in the design space. MOLT trains with FSDP2/EP/CP on NVIDIA AutoModel; slime's FSDP backend is experimental; MOLT and OpenRLHF orchestrate vLLM under Ray. $^1$RL code counts every Python file used by the RL entry path, trainer, rollout, orchestration, experience/advantage/loss, and imported model/utility/parallelism code, excluding pure SFT/DPO trainers, reward-model training, vendored code, tests, examples, and docs; counts trace the import graph from each RL entry point. Measured at verl 86e8123, slime 243773c, OpenRLHF b3d2927 (2026-06-16); MOLT 2026-07-07. LOC measures implementation footprint, not usability or correctness.

The five principles together define Molt’s readability‑first approach, turning a traditionally heavyweight RL stack into a lean, inspectable, and extensible research platform.

By codifying readability, minimalism, performance parity, algorithm‑aligned modularity, and strict correctness, Molt delivers a scalable yet approachable RL framework.

The Molt System Architecture

The runtime is a single asynchronous loop that ties agents, rollouts, and the policy actor.

MOLT turns the design principles into a concrete runtime: a single asynchronous loop that stitches together agents, rollout engines, and a policy actor. All data moves as token‑first records, guaranteeing that every training step uses exactly the tokens it generated.

**Figure 1.** The whole system: three components and one loop. MOLT composes user agents in plain Python, vLLM rollout engines behind a request router, and a single FSDP2 policy actor on NeMo AutoModel around one Ray asynchronous queue that implements the streaming pool and partial rollout. Trajectories flow token-exactly from the engines through the queue into training, and weight refit returns over NCCL directly to the engines, bypassing the request router.

The contract reduces the RL interface to a single Python module that receives token‑first observations and returns a reward, letting users write arbitrary Python logic while the framework handles token bookkeeping.

How does this Agent Contract differ from a typical OpenAI Gym environment?

Gym expects reset() and step() that return observations; the Agent Contract instead receives the full token trace and must return a Result with reward and termination. The surrounding framework supplies the observation loop and guarantees token‑first alignment, so the user never writes a reset function.

The framework calls

The agent’s

The loop proceeds to token 102, action “B”, and the agent again returns

At token 103, action “C” triggers the agent to return

The trainer receives the three token‑first records, each paired with its reward, and computes the policy loss without any additional alignment.

The contract lets arbitrary Python code decide the reward while the runtime guarantees that each reward is attached to the exact token that produced it.

Implementing the ChatAgent

Implements a plug‑in ChatAgent that runs via a loopback server, preserving token‑exact traces.

The Molt framework inserts a lightweight loopback server in front of any OpenAI or Anthropic SDK, turning ordinary API calls into a token‑exact, self‑contained chat session.

Think of the loopback server as a mail sorter: each outgoing request carries a stamped envelope (the session ID) that the sorter reads, records the exact letters (tokens) that go in and out, and then delivers the reply—all without ever leaving the postal system.

How does this ChatAgent differ from simply calling the OpenAI SDK directly?

Direct SDK calls send raw text over HTTP and lose the intermediate token boundaries; the loopback server intercepts the call, injects a session ID, and records each token as it is sent and received, guaranteeing a lossless token‑exact trace without any extra API flags.

Client sends the first request; the server records the inbound tokens [101, 102] and tags them with “sess‑1”.

Server returns the assistant’s tokens [201]; the TITO trace now contains (in → [101, 102], out → [201]).

Second user request arrives; inbound tokens [103, 104] are appended to the same session trace.

Assistant reply [202, 203] is captured; the full token‑exact trace for “sess‑1” is now {[101,102]→[201], [103,104]→[202,203]}.

The trace remains a single, contiguous token stream per session, so downstream reward models can evaluate the entire dialogue without any re‑tokenization step.

Minimal ChatAgent subclass and runner using Molt’s loopback server.

Because the same chat‑format dataset is rendered once by the server, both pure ChatAgent loops and environment‑driven agents consume identical inputs, eliminating data‑path confounds when swapping agent styles.

Scaling and Transport Consistency

How Molt scales to trillion‑parameter MoEs while keeping token transport exact and routing consistent.

Scaling MoE‑based RL agents traditionally forces a redesign of the serving stack whenever a new routing or parallelism scheme is introduced, and token‑level mismatches between rollout and training cause silent divergence.

Transport Consistency guarantees that every token that leaves the actor arrives at the rollout engine unchanged and that the same expert choices are replayed during training, eliminating hidden drift between inference and learning.

Actor sends token IDs [12, 45] to Engine A.

Engine A routes token 12 → expert E0, token 45 → expert E1 and returns expert indices [0, 1] alongside logits.

During training the actor receives the same indices [0, 1] and re‑executes the exact expert calls, producing identical forward activations.

Weight refit broadcasts the sharded parameter matrix; Engine A loads only its shard, others load theirs, no router involvement.

Because the expert indices are replayed verbatim, any tiny numerical difference in routing decisions cannot cause divergent sparse graphs, preserving exact consistency between rollout and training.

How does this differ from a typical token‑level API that also sends token IDs?

Typical APIs still invoke a tokenizer on the server side for each generation step, which can introduce nondeterministic tokenization boundaries. Transport Consistency removes that step entirely: the actor never hands a string to the engine, and the engine never returns a string, so the token stream is immutable from end‑to‑end.

With transport guaranteed, scaling MoEs becomes a matter of configuration: FSDP2 sharding, expert parallelism, and vLLM flags compose without code changes, turning a dense 4 B model launch script into a 1‑trillion‑parameter MoE launch by flipping a flag.

**Table 4.** **MoLT scaling knobs.** Actor-side parallelism composes with FSDP2 sharding; the vLLM side mirrors it per engine.

MoE consistency is a failure mode unique to sparse models: independent routing decisions in rollout and training can diverge due to floating‑point nondeterminism, causing the two sides to evaluate different sparse graphs.

Instead of trusting that two independent routers will pick identical experts, the rollout engine explicitly reports the chosen expert for each token, and the training actor replays that exact sequence.

Engine‑side features such as speculative decoding, prefix caching, and CUDA graphs are exposed as simple flags; unsupported flag combinations are rejected early by Molt’s guard, preventing silent misalignment.

All advantage estimators are plain Python functions that take rewards and group identifiers as inputs; selecting a different estimator is a one‑line flag change.

All loss terms share a single denominator: the total number of unmasked tokens in the optimizer‑step window, making the objective invariant to data‑parallel size or gradient‑accumulation depth.

For asynchronous rollouts, a per‑token importance weight corrects the mismatch between the behavior policy (rollout) and the target policy (training), while dynamic filtering removes degenerate groups.

Adding a new advantage estimator therefore reduces to writing a single function and exposing it via a flag—no subclassing, no engine changes.

Workflow and Resource Footprint

Shows how Molt’s lean design cuts code size and scales efficiently on large models.

Molt’s RL pipeline footprint is dramatically smaller than prior agents.

8.6K lines versus 62K for verl and 25K for slime (import‑graph counting, Tab. 1).

Throughput and memory were evaluated on the shipped Qwen3.6‑35B‑A3B recipe (a multimodal MoE policy) across two nodes (8 training + 8 rollout H100 GPUs). The benchmark probes three cost points—multi‑turn re‑prefill, generation, and actor memory—using a 2 K‑prompt DAPO‑Math dataset, 128 sequences per step, and bf16 precision.

The three‑step experiment workflow makes setting up, running, and observing large‑scale RL runs trivial.

Comparative Performance Benchmarks

MOLT matches slime’s throughput, confirming parity under a matched benchmark.

The modular design of Molt eliminates the monolithic bottleneck that forces custom engine forks for each new Agentic RL algorithm.

OpenRLHF is a lightweight interface that lets any reinforcement‑learning loop plug in a pretrained language model without rewriting the training code.

verl is a version‑controlled rollout layer that routes each conversation to the exact model snapshot that generated the preceding turn.

slime is the Megatron‑Core‑based reference stack that combines a 30 B‑parameter model with the SGLang serving engine.

MOLT and slime achieve statistically comparable optimizer‑step times, showing no clear superiority.

Mean step times are 119.4 ± 2.3 s (MOLT) versus 109.5 ± 10.3 s (slime) over three independent runs.

Related Frameworks and Approaches

We situate Molt among prior RL frameworks, async rollout systems, and trajectory capture approaches.

Molt executes a reference forward that runs fully asynchronously, allowing at most one policy‑version lag and enabling per‑token importance correction (Liu et al., 2025a). The benchmark reveals an upstream distributed‑MoE forward mismatch on a 128‑expert checkpoint, where actor log‑probabilities differ by roughly one nat and the [0.99, 1.01] gate rejects the batch, so measured step times reflect pure throughput without effective policy updates. In contrast, the 35B workload shows no such gap and the gate passes all sequences, leaving convergence‑parity validation for future work.

We position Molt against three families of systems: RL training frameworks, asynchronous rollout systems, and agent‑trajectory capture platforms. Its contribution relative to each family combines a readability‑first research surface, a token‑first agent contract, and a high‑performance PyTorch‑native implementation.

RL training frameworks such as HybridFlow/verl (Sheng et al., 2025b), OpenRLHF (Hu et al., 2024), NeMo‑RL, TRL, and Slime span a scale–complexity trade‑off: the former reach ultra‑large scales via broad multi‑backend surfaces or Megatron‑Core commitments, while the latter remain lightweight but do not target frontier‑scale agentic RL. Molt rejects this trade‑off by offering a single lean codebase, several times smaller than Megatron‑based stacks, yet capable of scaling from a quick single‑node experiment to a 700 B MoE on standard engines with comparable throughput.

Asynchronous and disaggregated RL systems include AReaL (Fu et al., 2025), StreamRL (Zhong et al., 2025), Laminar (Sheng et al., 2025a), DORA (Hu et al., 2026), RolloutPipe (Chen et al., 2026), Relax (Zhang et al., 2026b), and RollArt (Gao et al., 2025). These approaches add scheduler, staleness, or resource‑planning mechanisms; Molt shares disaggregation and completion‑driven streaming but does not claim those scheduler contributions, instead using a bounded persistent prompt‑group pool, generation‑time log‑probabilities, and per‑token correction that fit its readable core.

Training‑inference consistency work such as R3 (Ma et al., 2025) shows that divergent expert choices between rollout and training cause MoE instability, prompting replay of rollout routes during optimization. Molt implements this via vLLM’s native route capture and AutoModel RouterReplay, augmenting it with token‑level behavior probabilities and fail‑fast checks for unsupported feature combinations.

Agent‑trajectory and harness integration efforts such as Agent Lightning (Luo et al., 2025) and Polar (Xu et al., 2026) define standardized step‑granular trajectory protocols and black‑box harness proxies. Molt’s loopback server makes stock SDK calls token‑exact and segments trajectories under prefix‑rewriting compaction, addressing the first missing layer while leaving enterprise data‑conversion and unified control‑plane responsibilities to complementary systems.

Molt adopts existing RL objectives—REINFORCE++ (Hu et al., 2025), GRPO (Shao et al., 2024), DAPO (Yu et al., 2025), and GSPO (Zheng et al., 2025)—without proposing new ones, focusing its systems contribution on keeping the required quantities aligned, token‑exact, and consistently routed.

System Component Overview

Key components and comparative design points of Molt’s training stack.

This section visualizes Molt’s architecture and positions it against other agentic‑RL frameworks by enumerating their back‑ends, rollout engines, and code footprints.

Principle 1 (single‑call‑site advantage estimator) lets a developer add a new estimator by inserting a flag, a function, a metric, and a unit test—all visible in the same training loop iteration.

Evaluation asks whether Molt’s lean design (small RL code footprint, modular configuration) preserves throughput compared with a Megatron‑based stack under identical protocols.

Technical Appendix

Appendix details Molt’s scaling knobs and the reproduction setup.

Appendix A.1 lists Molt’s primary scaling knobs, grouped by system side. The actor side exposes tensor parallelism (‑fsdp.`tp_size`), expert parallelism (‑fsdp.`ep_size`) for up to 256 experts, context parallelism (‑fsdp.`cp_size`) for sequences longer than 32 K, and an optimizer‑offload flag (‑fsdp.offload optimizer) that moves Adam states to host memory while remaining checkpoint‑compatible. The vLLM rollout side mirrors these controls with per‑engine tensor parallelism (‑vllm.`tensor_parallel_size`), expert parallelism (‑vllm.`enable_expert_parallel`, where EP = TP × DP), data parallelism (‑vllm.`data_parallel_size`) that can raise EP beyond TP, a token‑budget scheduler (‑vllm.`max_num_batched_tokens`), and speculative decoding (‑vllm.`mtp_num_speculative_tokens`, disabled when set to 0).

Additional flags address MoE stability: ‑train.`routing_replay` enables rollout routing replay (the default in MoE recipes) and ‑actor.`freeze_moe_router` freezes the router as a coarser alternative to replay.

Appendix A.2 provides the reproduction note: the reported measurements use the shipped Qwen3.6‑35B‑A3B geo3k RL recipe on a 2‑node, 8‑GPU‑per‑node H100 cluster (8 training + 8 rollout GPUs) with 32 K context, CP8/EP8/TP1, four prompts each sampled four times, temperature 1.0, seq‑mask‑tis, and R3 enabled. The head‑to‑head benchmark follows Table 2, which pins the framework commits, and each configuration is released as a single‑command launch recipe with container specifications in the repository.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers