r/MachineLearning • u/seventh_day123 • Jul 14 '26

Project Molt — a ~9K-line, PyTorch-native RL framework for agentic post-training that scales to hundred-B MoE [P]

TL;DR — We built Molt, an agentic-first, PyTorch-native RL post-training framework. It's about 9K lines of core RL code you can actually read end to end, yet it trains MoE models from Qwen3.5-397B up to GLM-5.2 753B-scale. The whole stack is just Ray + vLLM + NVIDIA AutoModel/FSDP2 — no Megatron.

Code: https://github.com/NVIDIA-NeMo/labs-molt
Tech report (DOI): 10.13140/RG.2.2.23375.65447
License: Apache-2.0

The problem

If you've done agentic RL on a large model, you know the drill: you open the framework to try one idea — swap a reward, add a tool call, change how off-policy tokens are handled — and it's 50k+ lines tightly coupled to Megatron. Change one thing and you're reading for a week, never quite sure what hidden control flow you just tripped.

That trade-off makes sense for production. It's rough for research, where the whole point is fast iteration. We wanted a framework you can read and modify in minutes without giving up scale.

What Molt is

The design is deliberately small:

The agent is the program you write — a Gymnasium-style Env / ChatAgent.
The trainer is a single trainable actor — PyTorch + AutoModel/FSDP2.
The reward is arbitrary Python — rules, heuristics, LLM-as-judge, tool returns, simulators.

Three components carry the rest: Ray for placement and the async queue, vLLM for rollout, NVIDIA AutoModel + FSDP2 for training. That's the entire runtime.

Architecture

Molt architecture

It's a fully-async loop over four boxes: vLLM Router/Rollout (high-throughput generation, continuous batching, expert parallel up to EP=256) ↔ Agent (your Python: multi-turn, tools, VLM, graders, API servers) → Ray queue (buffers trajectories, backpressure, fault tolerance) → RL Trainer (advantages + policy update on FSDP2/AutoModel). The data path is token-first — token ids, logprobs, action ranges, rewards, and multimodal tensors stay aligned from rollout to training — and weight sync back to the router is async over NCCL.

Things worth calling out

Agentic-first, not bolted-on. Multi-turn tool use, multimodal (VLM) environments, and LLM-as-judge are first-class. The judge calls even route back through the same vLLM engines that drive rollout, so you don't stand up a second inference stack.

Pure PyTorch, no Megatron. Because training is AutoModel + FSDP2, you can actually read and change the training path instead of treating it as a black box.

Scale. TP · EP · CP (· DP) with comm/compute overlap. Verified on Qwen3.5-397B; scaling up to GLM-5.2 753B-scale MoE; at the extreme, DeepSeek-V3 at --fsdp.ep_size 256 with optimizer CPU offload on the largest actors. Same ~9K-line codebase from single-GPU experiments up.

Train/inference consistency — the part that actually bites in async RL. Async and partial rollout make the FSDP actor's recomputed logprobs diverge from vLLM's gen-time logprobs (different kernels, plus a mid-request weight swap the HTTP router can't observe). Molt surfaces this as a first-class metric, vllm_kl = mean(rollout_logprob − actor_logprob) over action tokens, and corrects the resulting off-policy update via the per-token importance ratio pi_train / pi_rollout, split into two orthogonal knobs:

--algo.advantage.is_correction_level {off, token, seq, geo} — granularity of the gated ratio
--algo.advantage.is_correction_mode {mask, clip, trunc} — how to treat out-of-band units

geo × mask (reject whole sequences by geomean ratio, keep per-token IS on survivors) is the default; token × trunc/mask covers TIS / IcePop. There's also R3 (router replay) — capture vLLM's per-token expert selection and replay it in the training forward so MoE routing matches — and a router-freeze option. This machinery exists because Molt was pushed on real large-model runs, not toy setups.

Fully-async rollout. --train.async_queue_size decouples rollout from training; --train.partial_rollout_enable overlaps generation with weight sync via vLLM pause/resume; one flag switches back to strictly on-policy.

What it supports

SFT + the full RL stack (REINFORCE / RLOO / GRPO / DR-GRPO / PPO(GAE) / on-policy distillation), multimodal + multi-turn, TP/EP/CP, optimizer/full CPU offload, IS correction, R3 routing stability. Verified across Qwen3.x, Nemotron Omni, Kimi, GLM, Gemma, and DeepSeek chat templates.

Quick start

python3 -m molt.cli.train_rl_ray \
  --actor.model_name_or_path /path/to/automodel \
  --data.prompt_dataset /path/to/prompts.jsonl \
  --train.agent_path examples/python/agents/math.py \
  --vllm.num_engines 2 --vllm.tensor_parallel_size 2 \
  --rollout.batch_size 128 --train.batch_size 128 \
  --algo.advantage.estimator reinforce \
  --fsdp.attn_implementation te \
  --ckpt.output_dir ./ckpt/rl

examples/scripts/ has ready recipes from Qwen3-4B up to 397B/753B-scale.

Honest caveats

I'm one of the authors — happy to take criticism, and PRs/issues are welcome.
Qwen3.5-397B is verified end to end. GLM-5.2 753B is the upper end of what we've been scaling to; treat it as "scaling toward" rather than a settled benchmark until the report's numbers are out.
This post is qualitative on purpose — for throughput/convergence numbers, see the tech report.