labs-molt

Molt is agentic-first and PyTorch-native. The agent is the program; the trainer is a single actor; reward is any Python you write inside an Env or ChatAgent — graders, multi-turn tools, VLM environments, LLM-as-judge. Three components carry the rest — Ray for placement and async queues, vLLM for rollout, NVIDIA AutoModel + FSDP2 for training in pure PyTorch. That is the whole stack: ~9.2K lines of RL code that scale to 1T-class MoE on vLLM with TP / EP / CP — think DeepSeek-V3 at --fsdp.ep_size 256, Adam CPU offload for the largest actors. One agent API, one trainable actor, clean enough to read end-to-end.

News

• 2026-09 · Molt now supports FlashREINFORCE, critic-free single-rollout RL with stable training beyond 6,000 steps — see the quick start.

Architecture

Three boxes. One async loop.

Ray owns placement and the async queue between the three boxes — that is the entire runtime. The contract is token-first: token ids, logprobs, action ranges, rewards, and multimodal tensors stay aligned from rollout to training. Anything you can compute in Python is a valid reward, including LLM-as-judge calls back through the same vLLM engines that drive rollout.

✨ Why Molt

	What you get	Why it matters for research
Agentic-first	One Gymnasium-aligned API — Env.step() or ChatAgent.run() — covers graders, multi-turn tools, VLM environments, and OpenAI/Anthropic-compatible servers	The agent is the program — iterate on environments in plain Python, the trainer stays untouched
⚙️ Fully-async runtime	Ray placement, async rollout queues, vLLM engines, partial rollout, weight sync	Rollout, training, and weight sync overlap — a DeepSeek-V3-class actor stays fed without bespoke infra
PyTorch-native, AutoModel-first	FSDP2 + NVIDIA AutoModel, pure PyTorch end-to-end	Hack the model in the language you already write; no backend ceremony
Single-actor simplicity	One actor, optional KL reference — the whole RL graph fits on a page	Every gradient is explicit; every loss term is one file away
Frontier-scale MoE	AutoModel + FSDP2 + TP / EP / CP + Adam CPU offload, MoE-native — e.g. DeepSeek-V3 with --fsdp.ep_size 256	The same script that trains 8B scales to 1T-class MoE — no rewrite between scales
Token-first contract	Aligned token ids, logprobs, action ranges, rewards, multimodal tensors	Multi-turn, VLM, and tool-call traces share one format end-to-end
Small, hackable surface	~9.2K LOC of RL code across 3 thin layers	Fork one layer without touching the others — read it in an afternoon

How It Compares

The RL ecosystem optimizes for breadth. Molt optimizes for agentic research velocity at scale — the smallest PyTorch-native stack that still drives fully-async agentic RL at frontier MoE scale on vLLM.

	** Molt**	OpenRLHF	verl	slime
Training backend	PyTorch / FSDP2 + NVIDIA AutoModel	DeepSpeed ZeRO-3	FSDP / FSDP2 / Megatron	Megatron (FSDP exp.)
Rollout engine	vLLM (Ray)	vLLM (Ray)	vLLM / SGLang / TRT-LLM	SGLang only
RL topology	actor (+ optional PPO critic)	actor + critic + RM	actor + critic + RM	actor + critic + RM
Reward source	agent Python	agent / endpoint / RM	agent / RM / endpoint	rollout fn / RM
Parallelism	TP / EP / CP, MoE-native	ZeRO-3 / FSDP	TP / PP / EP / SP	TP / PP / DP / CP / EP
Multimodal	VLM RL, multi-turn tool calls	VLM RL (v0.10+)	Qwen2.5-VL, Kimi-VL	geo3k VLM
Config surface	CLI flags only	CLI + scripts	Hydra + YAML	CLI + YAML
RL code size¹	~9.2K LOC	~7.2K	~62K	~25K
Design center	agentic-first research	RLHF coverage	production breadth	Megatron throughput

One framework, one job. Molt is the smallest PyTorch-native stack that takes an NVIDIA AutoModel from SFT to frontier-scale agentic RL on vLLM. Read every line that touches your gradients, in plain PyTorch.

¹ RL code = every Python file the framework's RL path uses — online trainer, rollout, Ray orchestration, experience/advantage/reward/KL/loss, actor/critic/RM inference, plus shared models, utils, parallelism, and kernels the RL training command depends on. Excludes pure SFT, DPO/KTO/IPO trainers, reward-model training, distillation, vendored third-party code, tests, examples, scripts, and docs. Counts code lines only (blank and comment-only lines excluded). Measured by tracing the import graph from each RL entry point (molt.cli.train_rl_ray, openrlhf.cli.train_ppo_ray, verl.trainer.main_ppo); slime loads its Megatron/SGLang backends lazily, so its core slime/ package plus its slime_plugins/ model-zoo (+~4.7K — the in-repo model code its RL path uses, counted on the same basis as molt's models/) are counted, minus SFT/distillation. Molt measured 2026-07-20 on this repo; the others measured 2026-06-16 at each repo's then-latest main HEAD (verl 86e8123, slime 243773c, OpenRLHF b3d2927).

Supported Scope

Training & runtime

Area	Support
SFT	molt.cli.train_sft
RL	vLLM-backed online RL via molt.cli.train_rl_ray
Runtime	Ray placement, async rollout queues, vLLM engines, partial rollout sync
Model scale	AutoModel + FSDP2 with TP / EP / CP, MoE-native — e.g. DeepSeek-V3 at --fsdp.ep_size 256
Model backend	NVIDIA AutoModel is the primary path — native CP / EP / TP, custom MoE+EP parallelizer, TE fused attention; everything model-side aligns with AutoModel's own recipes. The HF transformers path is a non-preferred fallback (AutoModel drops to it only when a model has no native class) supporting text + flash_attention_2 + packing only — no CP / EP / TP
Optimizer	adam (default), with CPU offload for the largest actors (--fsdp.offload optimizer). muon (Newton–Schulz via Dion: Muon for 2D weights and grouped MoE experts, AdamW for embeddings / head / norms) is experimental — runs distributed (FSDP / EP) but has shown no consistent win over adam yet, which stays the recommended default

Agents & rewards

Area	Support
Agent interface	--train.agent_path with Env or ChatAgent subclass + an AgentRunner
Reward source	Result(reward=...) returned from Env.step or ChatAgent.run
Modalities	Text and VLM prompts, including image payloads
Chat templates	Assistant spans (SFT loss mask + multi-turn rollout stitching) are derived from the model's own chat template — no hard-coded markers. Verified on ChatML (Qwen3.x, Nemotron omni3), Kimi-K2.6, GLM, Gemma and DeepSeek

Algorithms

Area	Support
Estimators	reinforce, reinforce_baseline, rloo, grpo, dr_grpo, gae (PPO), on_policy_distill
PPO critic	--algo.advantage.estimator gae adds a value model: its own Ray group (CriticModelActor), colocated on the actor's GPUs by default or disaggregatable, GAE advantages (--algo.advantage.lam) + clipped value loss (--critic.value_clip), own optimizer/LR (--critic.adam.lr) and resumable _critic checkpoint. Built on NeMoAutoModelForCausalLM + a scalar value head, so it keeps the native TP / EP / CP path
Distillation	On-policy distillation — per-token reverse KL to a frozen teacher, via --algo.advantage.estimator on_policy_distill + --ref.model_name_or_path
IS correction	Train/rollout logprob-mismatch correction for off-policy / async rollout: is_correction_level {off,token,seq,geo} × is_correction_mode {mask,clip,trunc} (covers TIS, IcePop, seq-mask-tis; see IS correction below)
KL	Optional reference workers when --algo.kl.init_coef > 0 (the reference doubles as the distillation teacher)

MoE routing stability

Area	Support
Router replay (R3)	--train.routing_replay — vLLM's per-token top-k selection replayed in the training forward; details in the MoE routing stability section under Scaling Knobs
Router freeze	--actor.freeze_moe_router holds the gate/router weights fixed so vLLM and the actor keep routing tokens to the same experts. Stabilizes MoE RL / distillation and shrinks the same rollout-vs-train logprob gap the IS-correction filters address — a router that drifts between refits is a large source of that gap

IS correction (train/rollout logprob mismatch)

Async and partial rollout make the FSDP actor's recomputed pi_train diverge from vLLM's gen-time pi_rollout (different kernels, plus a mid-request weight swap the HTTP router can't observe). Molt corrects the resulting off-policy update with the per-token importance ratio pi_train / pi_rollout, gated by two knobs:

• --algo.advantage.is_correction_level {off, token, seq, geo} — granularity of the gated ratio. off disables correction; token gates each token's own ratio; seq/geo aggregate a sequence's ratios (exp(sum) / exp(mean)) into a per-sequence rejection filter (kept sequences still carry their per-token IS weight), so they require mode mask.
• --algo.advantage.is_correction_mode {mask, clip, trunc} — treatment of a unit outside the band. mask drops it (zero gradient); clip clamps its weight into the band; trunc clamps only the upper tail.
• --algo.advantage.is_correction_threshold [LOW] HIGH — the [low, high] band on the ratio (recipes use a tight 0.99 1.01); a single value is an upper bound only.

The named schemes and their prior art:

Flags	Scheme	Prior art
level token mode trunc	truncated IS	TIS
level token mode mask	token masking	IcePop
level token mode clip	token clip	per-token weight clamp
level geo mode mask	seq-mask-tis (recipe default)	MIS-style sequence masked importance sampling
level seq mode mask	product-ratio reject	sequence log-ratio sum

References: TIS (truncated importance sampling of the train/infer ratio), IcePop (token-level masking of out-of-band ratios), and MIS (masked importance sampling, Yingru Li — sequence-level masked IS, which motivates the seq/geo rejection filter).

Installation

First clone the repo — the launch scripts, agents, and recipes live here, and examples/scripts/docker_run.sh mounts this checkout into the container. For local (non-container) development, add the editable install: it pulls the exact git-pinned AutoModel this repo is validated against, so R3 routing replay and Muon work out of the box:

git clone https://github.com/NVIDIA-NeMo/labs-molt.git
cd labs-molt
pip install -e ".[vllm]"          # local development only — the container bakes everything in

Requires CUDA 13. The git-pinned AutoModel is only compatible with the CUDA-13 torch build (torch==2.13.0+cu130); a CUDA-12 environment will not work. If your host driver is older than 580 (native CUDA 13), use the container instead — it ships the CUDA forward-compatibility layer and enables it automatically.

The recommended path is the project container (dockerfile/Dockerfile). It bakes the full CUDA-13 stack — torch 2.13 · vLLM · TransformerEngine · flash-attn · mamba · DeepEP · NVIDIA AutoModel — built for A100 / H100 / H200 / B200·GB200, so it runs SFT and RL as-is with no local dependency wrangling. Pull the prebuilt image from Docker Hub:

docker pull hijkzzz/molt:latest   # or a pinned release: hijkzzz/molt:0.1.7

...or build it yourself from the Dockerfile (e.g. to change the CUDA / vLLM / AutoModel pins):

docker build -f dockerfile/Dockerfile -t hijkzzz/molt:latest .

From 0.1.9 the Docker Hub tags are multi-arch (amd64 + arm64), so the same `docker pul