Metadata-Version: 2.4
Name: molt-rl
Version: 0.1.12
Summary: A simple Ray + vLLM + FSDP2 (AutoTP/EP/CP) stack for SFT and RL.
Author: NVIDIA CORPORATION & AFFILIATES
License: Apache-2.0
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Environment :: GPU :: NVIDIA CUDA
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: System :: Distributed Computing
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: NOTICE
Requires-Dist: accelerate
Requires-Dist: aiohttp
Requires-Dist: datasets
Requires-Dist: einops
Requires-Dist: grpcio>=1.74.0
Requires-Dist: huggingface_hub>=1.0.0
Requires-Dist: jsonlines
Requires-Dist: optree>=0.15.0
Requires-Dist: packaging
Requires-Dist: peft
Requires-Dist: pylatexenc
Requires-Dist: pynvml>=12.0.0
Requires-Dist: ray[default]==2.59.0
Requires-Dist: sympy
Requires-Dist: tensorboard
Requires-Dist: tokenizers<0.24.0,>=0.23.1
Requires-Dist: torch
Requires-Dist: torchdata
Requires-Dist: torchmetrics
Requires-Dist: tqdm
Requires-Dist: transformers==5.18.0
Requires-Dist: vllm
Requires-Dist: vllm-router>=0.1.15
Requires-Dist: wandb
Requires-Dist: wheel
Requires-Dist: nemo-automodel>=0.5.0
Provides-Extra: vllm
Requires-Dist: vllm==0.31.0; extra == "vllm"
Provides-Extra: vllm-latest
Requires-Dist: vllm>=0.24.0; extra == "vllm-latest"
Provides-Extra: flash-attn-2
Requires-Dist: flash-attn==2.8.3; extra == "flash-attn-2"
Dynamic: author
Dynamic: classifier
Dynamic: description
Dynamic: description-content-type
Dynamic: license
Dynamic: license-file
Dynamic: provides-extra
Dynamic: requires-dist
Dynamic: requires-python
Dynamic: summary

<div align="center">

# 🦋 Molt

**An agentic-first RL framework for research.**

Ray · vLLM · NVIDIA AutoModel — the smallest PyTorch / HuggingFace-native stack for
1T-class fully-async, multimodal, multi-turn agentic RL.

<br/>

[![License](https://img.shields.io/badge/License-Apache_2.0-2563eb?style=flat-square)](LICENSE)
![Python](https://img.shields.io/badge/Python-3.10+-3776AB?style=flat-square&logo=python&logoColor=white)
![PyTorch](https://img.shields.io/badge/PyTorch-native-EE4C2C?style=flat-square&logo=pytorch&logoColor=white)
![NVIDIA AutoModel](https://img.shields.io/badge/Training-NVIDIA_AutoModel-76B900?style=flat-square&logo=nvidia&logoColor=white)
![vLLM](https://img.shields.io/badge/Rollout-vLLM-7c3aed?style=flat-square)
![Ray](https://img.shields.io/badge/Runtime-Ray-028CF0?style=flat-square)
![RL code](https://img.shields.io/badge/RL_code-~9.2K_LOC-10b981?style=flat-square)
[![Tech Report](https://img.shields.io/badge/Tech_Report-arXiv:2607.21653-b31b1b?style=flat-square)](https://www.researchgate.net/publication/409325071_Molt_A_Scalable_PyTorch-Native_Training_Framework_for_Agentic_Reinforcement_Learning)
<a href="https://deepwiki.com/NVIDIA-NeMo/labs-molt"><img src="https://devin.ai/assets/deepwiki-badge.png" alt="Ask DeepWiki.com" style="height:20px;"></a>

<br/>

[**Paper**](https://arxiv.org/abs/2607.21653) ·
[**News**](#-news) ·
[**Architecture**](#-architecture) ·
[**Why Molt**](#-why-molt) ·
[**Installation**](#-installation) ·
[**Quick Start**](#-quick-start) ·
[**Scope**](#-supported-scope) ·
[**Agent Contract**](#-agent-contract) ·
[**Recipes**](#-recipes) ·
[**Scaling**](#-scaling-knobs) ·
[**Deep dives**](#-deep-dives) ·
[**Contributing**](#-contributing)

| Package | SFT | RL | Runtime |
|---|---|---|---|
| `molt` | `molt.cli.train_sft` | `molt.cli.train_rl_ray` | vLLM |

</div>

Molt is **agentic-first** and **PyTorch / HuggingFace-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](https://www.alphaxiv.org/abs/2609.flashreinforce-asynchronous-rl-agentic-models), critic-free single-rollout RL with stable training beyond 6,000 steps — see the [quick start](examples/scripts/quick_start/rl_flash_reinforce_r1d_1p5b.sh).

## 🧩 Architecture

Three boxes. One async loop.

<p align="center">
  <img src="assets/molt.jpg" alt="Molt architecture: Agent · vLLM rollout · Ray async queue · single-actor AutoModel/FSDP2 trainer, fully async" width="920"/>
</p>

**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 |

## 📦 Installation

Clone the repo: the launch scripts, agents and recipes live here, and the container mounts this checkout.

```bash
git clone https://github.com/NVIDIA-NeMo/labs-molt.git
cd labs-molt
```

### 🐳 Container (recommended)

`dockerfile/Dockerfile` bakes the whole CUDA-13 stack (torch 2.13, vLLM, TransformerEngine, flash-attn,
mamba, DeepEP, NVIDIA AutoModel) for A100 / H100 / H200 / B200 / GB200. SFT and RL run in it as-is.

```bash
docker pull hijkzzz/molt:latest                                 # or a pinned release: hijkzzz/molt:0.1.10
docker build -f dockerfile/Dockerfile -t hijkzzz/molt:latest .  # to change the CUDA / vLLM / AutoModel pins
```

Tags from 0.1.9 are multi-arch (amd64 + arm64): the same pull works on x86 and on Grace-Blackwell hosts,
and the one Dockerfile builds both (`docker buildx build --platform linux/amd64,linux/arm64 --push` to
publish a tag).

### 💻 Local install

For development outside the container. It pulls the exact git-pinned AutoModel this repo is validated
against (`setup.py`, `AUTOMODEL`), so R3 routing replay and Muon work out of the box.

```bash
pip install -e ".[vllm]"
```

Requires CUDA 13 (`torch==2.13.0+cu130`); a CUDA-12 environment will not work. With a host driver older
than 580, use the container: it ships the CUDA forward-compatibility layer and enables it automatically.

The `molt-rl` package on PyPI cannot carry the git-pinned AutoModel, so it lags the pin and is not a
supported install path right now. Use the container or the local install.

### 🪶 Optional backend: AutoModel-Slim

**The default is upstream NVIDIA AutoModel**, pinned in `setup.py`; every command above installs it, and
nothing in this subsection applies unless you opt in.

[`automodel-slim`](https://github.com/NVIDIA-NeMo/labs-molt/tree/automodel-slim) is molt's own copy of that
pinned AutoModel commit, trimmed to what molt uses and maintained in this repo. Like upstream it is
PyTorch-native and Hugging Face-native: Hugging Face checkpoints in, native FSDP2 implementations for
training, Hugging Face safetensors out for vLLM and `transformers`. Same package name (`nemo_automodel`),
same import paths, same API, so molt's code does not change when you switch.

It keeps the model families molt trains (Qwen2 / 2.5, Qwen3 / 3.5 / 3.6 / 3.8, DeepSeek V4.1 Flash,
GLM 5.x, Gemma 4, Nemotron 3, Muse Glimmer, Inkling) and the training stack behind them (FSDP2 with
TP / EP / CP, TransformerEngine attention with THD packing, DeepEP / HybridEP MoE dispatch, router replay,
LoRA, FP8, Dion / Muon, DCP checkpoints with consolidated HF export); everything else is removed, 758 files /
282k lines down to 252 / ~97k. Forward logits are bit-exact with upstream on every kept family, molt's RL
e2e metrics match, and the branch runs its own seven-minute CI on molt's image and runners. The per-family
parity table, the kept / removed lists and the maintenance rules are in the branch's README.

```bash
MOLT_AUTOMODEL=slim pip install -e ".[vllm]"                                          # local install
docker build --build-arg MOLT_AUTOMODEL=slim -f dockerfile/Dockerfile -t molt:slim .   # image
EXTRA_PYTHONPATH=/path/to/automodel-slim ...                                          # slurm recipes: a checkout wins over the baked-in package
```

`setup.py` holds both pins (`AUTOMODEL`); the slim install follows the branch head, so a merge there is
live on the next image build or reinstall. Model and parallel-stack changes go to that branch as PRs, with
the same `cicd` label check as this repo. Request reviews with `/review`; legacy
`/claude review` comments only receive migration guidance.

## 🚀 Quick Start

### 📘 SFT

```bash
torchrun --standalone --nproc_per_node=8 -m molt.cli.train_sft \
  --model.model_name_or_path /path/to/automodel \
  --data.dataset /path/to/sft.jsonl \
  --data.input_key input \
  --data.output_key output \
  --ckpt.output_dir ./ckpt/sft \
  --fsdp.attn_implementation te
```

SFT uses the same AutoModel/FSDP2 model-loading path as RL.

### 🎮 RL

```bash
python3 -m molt.cli.train_rl_ray \
  --actor.model_name_or_path /path/to/automodel \
  --data.prompt_dataset /path/to/prompts.jsonl \
  --data.input_key input \
  --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 \
  --train.micro_batch_size 1 \
  --algo.advantage.estimator reinforce \
  --algo.kl.init_coef 0 \
  --fsdp.attn_implementation te \
  --ckpt.output_dir ./ckpt/rl
```

Common RL switches:

| Goal | Flags |
|---|---|
| Disable reference workers | `--algo.kl.init_coef 0` |
| Enable KL regularization | Set `--algo.kl.init_coef` above zero and place reference workers with `--ref.num_nodes`, `--ref.num_gpus_per_node`, or `--train.colocate_fsdp_models` |
| Compare samples per prompt | `--rollout.n_samples_per_prompt 8` plus `reinforce_baseline`, `rloo`, `grpo`, or `dr_grpo` |
| Weigh every prompt the same in the loss | `--actor.loss_agg_mode prompt-mean-token-mean` (token mean inside each prompt's rollouts, then mean over prompts; `seq-mean-token-mean` weighs every rollout the same, the default `token-mean` every token) |
| Decouple rollout and training | `--train.async_queue_size 2` |
| Keep rollout alive during sync | `--train.partial_rollout_enable` |
| Filter by agent scores | `--algo.dynamic_filtering_enable --algo.dynamic_filtering_range 0.0 1.0` |
| Correct async rollout logprobs | `--algo.advantage.is_correction_level geo` (seq-mask-tis; token-level adds `--algo.advantage.is_correction_mode clip/trunc/mask`) |
| FlashREINFORCE (critic-free, single-rollout) | `--train.force_on_policy --algo.advantage.estimator flash_reinforce --algo.advantage.is_correction_level seq --algo.advantage.is_correction_gating binary_kl --algo.advantage.is_correction_threshold 5e-3 --actor.loss_agg_mode seq-mean-token-mean` — see [the quick start](examples/scripts/quick_start/rl_flash_reinforce_r1d_1p5b.sh) |
| Score centering (off-policy REINFORCE without IS) | `--algo.advantage.score_centering_top_k 32`: records the rollout sampler's top-k log-probs per token and subtracts its expected score from the policy loss ([paper](https://arxiv.org/abs/2609.20807)). Use it **instead of** an IS gate (`is_correction_level off`), or compose it with a token-level weight only: `--algo.advantage.is_correction_level token --algo.advantage.is_correction_mode trunc` (the paper's TIS variant) or `mode mask` (IcePop-style MIS). It cannot be stacked on the sequence-level gate the FlashREINFORCE recipe uses (`is_correction_level seq/geo`): a sequence gate has no per-candidate form, and its surviving tokens keep the raw ratio, under which the centering term is identically zero |
| Freeze MoE routing (stabilize MoE RL) | `--actor.freeze_moe_router` |
| On-policy distillation | `--algo.advantage.estimator on_policy_distill --ref.model_name_or_path /path/to/teacher` |
| Independent eval sampling | `--eval.temperature`, `--eval.top_p`, `--eval.max_new_tokens`, `--eval.n_samples_per_prompt` (unset ones fall back to rollout) |
| Eval a checkpoint (no training) | `--eval.eval_only --eval.dataset <path>` — score the eval set once and exit; vLLM holds the HF weights, so the policy/ref/critic FSDP actors are never built and their GPUs go to the eval |
| Dump / replay a rollout batch | `--train.rollout_dump_dir <dir>` then `--train.rollout_replay_dir <dir>` re-runs training on it without regenerating |
| Check weight-update coverage | `--train.check_weight_update_equal` warns which vLLM params a broadcast left stale |

## 📊 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.

<details>
<summary>¹ How the RL-code line counts were measured</summary>

> ¹ 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`).

</details>

## 🎯 Supported Scope
### ⚙️ Training & runtime
- **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-Omni), Kimi-K2.6, GLM, Gemma, and DeepSeek |
### 🧮 Algorithms
- **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 |

Deep dives on IS correction, router replay, MTP rollout and LoRA are collected [below](#-deep-dives).

## 🤖 Agent Contract

Every RL run points at one Python module:

```bash
--train.agent_path /path/to/agent.py
```

The module must export `AgentRunner`. Choose **one** of two paths:

### 🧭 1. `Env` — framework owns the LLM loop *(Gymnasium-style step/reset)*

```python
from molt.agents import Env, Result, StepEnvRunner


class MathEnv(Env):
    async def step(self, state) -> Result:
        # state: observation_text, action_text, label, sampling_params
        reward = grade(state["action_text"], state["label"])
        return Result(reward=reward, terminated=True)


class AgentRunner(StepEnvRunner):
    def __init__(self):
        super().__init__(MathEnv)
```

The framework drives vLLM, tokenization, multimodal accounting, and
per-turn budgets. Your `step()` returns a `Result`; the framework chains
turns until `terminated` or `truncated`.

### 💬 2. `ChatAgent` — you own the loop via the OpenAI **or** Anthropic SDK

<details>
<summary>Full ChatAgent example</summary>

```python
from openai import AsyncOpenAI
from molt.agents import ChatAgent, ChatAgentRunner, ChatContext, Result


class MyAgent(ChatAgent):
    async def run(self, ctx: ChatContext) -> Result:
        # ctx.base_url carries the session id and auto-captures the token
        # trace — no extra_body, no logprobs=True, no session plumbing.
        client = AsyncOpenAI(base_url=ctx.base_url, api_key=ctx.api_key)
        resp = await client.chat.completions.create(
            model=ctx.model_name,
            messages=[{"role": "user", "content": ctx.prompt}],
            max_tokens=ctx.sampling_params.max_tokens,
            temperature=ctx.sampling_params.temperature,
        )
        return Result(reward=grade(resp.choices[0].message.content, ctx.label))


class AgentRunner(ChatAgentRunner):
    def __init__(self):
        super().__init__(MyAgent)
```

</details>

A multi-turn agent that stops on its own turn cap should return
`Result(truncated=True)` (see `examples/python/agents/chat_geo3k.py`); the
server marks generation-length / context truncation by itself.

The same server speaks the Anthropic wire too — point `AsyncAnthropic` at
`ctx.session_url` (the session root *without* `/v1`; the SDK appends
`/v1/messages` itself); everything else is identical:

```python
from anthropic import AsyncAnthropic

client = AsyncAnthropic(base_url=ctx.session_url, api_key=ctx.api_key)
msg = await client.messages.create(
    model=ctx.model_name,
    messages=[{"role": "user", "content": ctx.prompt}],
    max_tokens=ctx.sampling_params.max_tokens,
)
text = msg.content[0].text
```

A FastAPI vLLM server is auto-launched on loopback under the session URL.
External HTTP callers (browser automation, eval harnesses, OSWorld, …) hit the
same engine through either `/v1/chat/completions` (OpenAI) or `/v1/messages`
(Anthropic) — both decode to one token-exact accumulation.

#### Context compaction → multiple step-samples per rollout

Each chat call carries the prior turn's **exact** tokens forward and appends only
the new delta, so a multi-turn episode stitches into one monotonic token-exact
trajectory. But a long-horizon agent often **compacts** its context — summarizing
or dropping old turns to stay under the window (e.g. a `/compact` step) — which
*rewrites* the prefix, so it's no longer a clean extension of what was tokenized.
The model's own chat template can rewrite the prefix too: Qwen3-style templates
re-render a prior assistant turn without its `<think>` block once a newer user
query follows.

The server detects this automatically: when an incoming request rewrites the
prefix instead of extending it, it **seals the current segment and starts a fresh
token-exact segment** from the re-templated conversation. One
rollout therefore emits several segment trajectories — they share the rollout's
reward and `rollout_id`, so group baselines (GRPO/RLOO/…) deduplicate them to *one
reward per rollout* while each segment still contributes its own generated tokens
to the policy gradient (the same step-sample contract multi-turn agents use). No
agent-side change is needed — it works on both wires, including external harnesses
(Claude Code, opencode, AgentScope, …) whose compaction is opaque to us.

### 📋 `Result` fields

| Field | Meaning |
|---|---|
| `reward` | Scalar reward consumed by the trainer (required) |
| `observation` | Next-turn observation text (multi-turn only) |
| `terminated` | Episode finished naturally; defaults to `True` |
| `truncated` | Cut off externally (max turns, length, etc.) |
| `info` | Optional dict of scalar diagnostics for logging |
| `score` | Optional dynamic-filtering / dashboard score (defaults to `reward`) |
| `images` | Optional list of next-turn images |
| `sampling_params` | Optional per-turn override |

Four reference agents ship under `examples/python/agents/`:

```bash
--train.agent_path examples/python/agents/math.py          # Env: single-turn boxed grader
--train.agent_path examples/python/agents/geo3k.py         # Env: VLM multi-turn + Python tool
--train.agent_path examples/python/agents/chat_minimal.py  # ChatAgent: hello-world chat loop
--train.agent_path examples/python/agents/chat_geo3k.py    # ChatAgent: VLM multi-turn + Python tool
```

Two optional length penalties, both off by default and **mutually exclusive** -- they are
two answers to the same question, a rollout that runs out of context trips either one, and
the CLI rejects enabling both. Each shapes the merged per-rollout reward before
`--reward.clip_range` (group estimators merge a rollout's segments first; REINFORCE and GAE
shape each sample separately):

- `--reward.overlong_buffer_len B` ramps a DAPO-style penalty up to
  `--reward.overlong_penalty_factor` (default 1) over the last `B` tokens of `--data.max_len`.
- `--reward.stop_properly_penalty_coef C` scales a truncated rollout's reward by `C` in
  `[0, 1]`, or replaces it with `C` when `C` is negative. Truncation covers the generation
  length limit, context exhaustion and the agent turn cap.

Both measure a rollout by its **peak segment**, never by a sum: history compaction resets the
context window, so segments of one rollout do not add up. The quantity is the context a
segment occupied -- prompt, generated tokens, tool feedback and the image tokens vLLM expands
the placeholders into -- rather than a generation count, so `--rollout.max_new_tokens` may
stay unset.

Raw `reward` metrics are preserved; whichever rule is on reports its mean correction as
`overlong_penalty` or `stop_properly_penalty`, so a nonzero value is also how you see that it
fired, and `rollout/truncated_rate` keeps counting truncation over the full generated set.

## 🍳 Recipes

Reference launch scripts live under `examples/scripts/`, all end-to-end
on the AutoModel + FSDP2 backend:

| Workflow | quick_start | slurm |
|---|---|---|
| Qwen3.6-35B-A3B VLM SFT on geo3k | `quick_start/sft_qwen3_6_35b.sh` | `slurm/sft_qwen3_6_35b.sh` |
| Qwen3.6-35B-A3B VLM RL on geo3k (multi-turn Python tool) | `quick_start/rl_qwen3_6_35b.sh` | `slurm/rl_qwen3_6_35b.sh` |
| Qwen3-4B dense SFT on text math | `quick_start/sft_qwen3_4b.sh` | `slurm/sft_qwen3_4b.sh` |
| Qwen3-4B dense RL on text math | `quick_start/rl_qwen3_4b.sh` | `slurm/rl_qwen3_4b.sh` |
| Nemotron-Omni-30B-A3B VLM RL on geo3k (hybrid SSM MoE, CP8+EP8) | — | `slurm/rl_omni3_30b.sh` |
| Nemotron-Omni-30B-A3B on-policy distillation | — | `slurm/rl_distill_omni3_30b.sh` |
| GLM-5.2 ~750B RL on text math (MLA + DSA sparse attention, EP256) | — | `slurm/rl_glm5_2_753b.sh` |

Quick-start single-node usage:

```bash
MODEL_PATH=/path/to/Qwen3-4B bash examples/scripts/quick_start/rl_qwen3_4b.sh
```

The geo3k VLM scripts (`rl_qwen3_6_35b.sh` / `sft_qwen3_6_35b.sh`) auto-prepare the
dataset on first run via `examples/python/utils/prepare_geo3k.py`. To pre-stage it
manually (or refresh it), run:

```bash
python3 examples/python/utils/prepare_geo3k.py --num-proc 8 --out-dir .tmp/geo3k
```

Or point `PROMPT_DATASET` / `EVAL_DATASET` at your own data.

Slurm usage:

```bash
# 1) SFT smoke on 2 interactive nodes
sbatch examples/scripts/slurm/sft_qwen3_6_35b.sh

# 2) RL smoke on 2 interactive nodes (auto-preps geo3k on first run)
sbatch examples/scripts/slurm/rl_qwen3_6_35b.sh

# 3) Scale RL to 4 nodes for convergence
sbatch --nodes=4 examples/scripts/slurm/rl_qwen3_6_35b.sh
```

### 🔧 Multi-turn Python tool env

`examples/python/agents/geo3k.py` is the VLM multi-turn recipe used by the
Qwen3.6 RL script. The model emits a `<tool_call>` invoking
`python_executor(code=...)`; the env runs the snippet in a sandboxed
subprocess and feeds the captured stdout back as a `<tool_response>` turn.
The loop runs up to `MAX_AGENT_TURNS` (agent default 5; the shipped Qwen3.6
recipes set 4 for the quick start and 10 on Slurm); the final
`<answer>ANSWER</answer>` (or `\boxed{ANSWER}` for legacy distributions) is
graded against the ground truth and becomes the reward.

### 🌐 OpenAI- / Anthropic-compatible server agent

For agents that already speak OpenAI Chat Completions or the Anthropic Messages
API, subclass `ChatAgent` (see `examples/python/agents/chat_minimal.py`). The
auto-launched server exposes both `/v1/chat/completions` and `/v1/messages`
against the rolling vLLM engines, so any external loop (browser automation, eval
harness, OSWorld, …) can drive the policy through a stock OpenAI or Anthropic SDK
— both wires decode to the same token-exact trajectory capture.

### 🎓 On-policy distillation

Distill a student toward a frozen teacher on the student's *own* on-policy
samples. A single switch —
`--algo.advantage.estimator on_policy_distill` — turns the reference model into
the teacher and makes the per-token **reverse KL** to it the entire training
signal: the advantage becomes `-kl_coef · (log π_student − log π_teacher)` with
no scalar reward, no group baseline, and no whitening, so the policy loss is the
policy-gradient estimator of the reverse-KL gradient that pulls the student
toward the teacher.

```bash
python3 -m molt.cli.train_rl_ray \
  --actor.model_name_or_path /path/to/student \
  --ref.model_name_or_path /path/to/teacher \
  --algo.advantage.estimator on_policy_distill \
  --data.prompt_dataset /path/to/prompts.jsonl \
  --data.input_key input \
  # vllm / fsdp / batch flags as in the RL quick start
```

Everything else is derived from the one switch, so pure distillation needs no
reward function and no task agent — only the teacher checkpoint. Selecting the
estimator forces `--algo.kl.estimator k1`, turns `--algo.kl.use_loss` off (the
KL flows through the advantage, not a separate loss term), defaults
`--algo.kl.init_coef` to `1.0`, and — when no `--train.agent_path` is given —
auto-selects a built-in single-turn, VLM-aware generator
(`molt/agents/distill_agent.py`) that samples one on-policy completion per
prompt and returns a dummy `0.0` reward the estimator ignores.

The teacher **must share the student's processor/tokenizer** so the per-token
logprobs align over the same (vision-expanded) sequence — typically a larger or
more-trained checkpoint from the same family. It loads inference-only and can be
colocated on the actor nodes (`--train.colocate_fsdp_models`) or given its own
`--ref.num_nodes`. Watch `kl` / `logprobs_diff` fall toward 0 as the student
matches the teacher; task accuracy is not the objective, so eval is off.

To distill a **multi-turn tool-use distribution** (matching how the student is
actually deployed), point `--train.agent_path` at the task's real agent (e.g.
`chat_geo3k.py`) — its reward is simply ignored by the estimator.
`examples/scripts/slurm/rl_distill_omni3_30b.sh` is a ready-made VLM example
built on the omni3 EP8 / CP8 / TE / DeepEP recipe.

## 🎛️ Scaling Knobs
Molt targets AutoModel custom models with FSDP2.

**Actor (FSDP2)**

| Mode | Flag | Note |
|---|---|---|
| Tensor parallel | `--fsdp.tp_size 2` | |
| Expert parallel | `--fsdp.ep_size 8` | `256` for DeepSeek-V3-class MoE |
| Context parallel | `--fsdp.cp_size 8` | 32K+ sequences; VLMs and MoE routing replay shard with the sequence |
| Optimizer CPU offload | `--fsdp.offload optimizer` | frees VRAM for the largest actors |

**vLLM rollout**

| Mode | Flag | Note |
|---|---|---|
| Tensor parallel | `--vllm.tensor_parallel_size 2` | |
| Expert parallel | `--vllm.enable_expert_parallel` | EP = TP × DP |
| Data parallel | `--vllm.data_parallel_size 4` | single-node mp; raises EP past TP (DeepSeek-V3-style TP8 + DP4 → EP32) |
| Scheduler token budget | `--vllm.max_num_batched_tokens 32768` | |
| MTP speculative decoding | `--vllm.mtp_num_speculative_tokens 1` | details in [Deep dives](#-deep-dives) |

**MoE stability**

| Mode | Flag | Note |
|---|---|---|
| Router replay (R3) | `--train.routing_replay` | details in [Deep dives](#-deep-dives) |
| Router freeze | `--actor.freeze_moe_router` | |

Context parallelism is delegated to AutoModel's `ContextParallelSharder`, so each model gets the sharding
its attention backend needs:

- round-robin for hybrid SSM / linear-attention models (Nemotron-Omni, Qwen3.5-MoE), flat THD streams for
  sparse-attention models (GLM-5.2 DSA);
- VLM vision towers and routing replay shard with the sequence, so `--fsdp.cp_size` composes with
  `--data.image_key` and `--train.routing_replay`;
- sample packing (`--fsdp.packing_samples`) is text-only and off by default; under CP it takes the THD path.

## 🔬 Deep dives

The long-form notes behind the knobs above.

### ⚖️ IS correction — train/rollout logprob mismatch

<details>
<summary>The knobs, the named schemes and their prior art</summary>

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 |
| `score_centering_top_k 32` (+ optional `level token mode trunc\|mask`) | score centering (no ratio; optionally on top of TIS / IcePop) | Score Centering |

References: **TIS** (truncated importance sampling of the train/infer ratio), **IcePop**
(token-level masking of out-of-band ratios), **MIS** (masked importance sampling, Yingru Li —
sequence-level masked IS, which motivates the `seq`/`geo` rejection filter), and
**Score Centering** ([arXiv 2609.20807](https://arxiv.org/abs/2609.20807)).

**Score centering** (`--algo.advantage.score_centering_top_k k`)

*The problem in one sentence.* The rollout engine (vLLM) and the trainer (FSDP) are two different
implementations of the same model, so they disagree slightly on every token's probability; plain
REINFORCE then nudges the trainer toward the engine's preferences at every token, even when the
reward is zero. That nudge is the drift.

*Notation.* Look at one position of one response. The engine picked token `y` from its
distribution `q`; the trainer assigns the same position its own distribution `p`; `v` runs over
the vocabulary; `∇` is the gradient with respect to the trainer's weights.

*Where the drift comes from.* REINFORCE's per-token update is `advantage × ∇log p_y`. Because `y`
was drawn from `q`, the average of `∇log p_y` over the engine's sampling is

```
E_{y~q}[∇log p_y]  =  Σ_v q_v ∇log p_v
```

If engine and trainer agreed (`q = p`) this is `Σ_v p_v ∇log p_v = Σ_v ∇p_v = ∇1 = 0`: with no
reward signal there is no net push. When they disagree it is not zero, so every update carries a
reward-independent push — the drift.

*What score centering does.* It subtracts the average push before the advantage is applied:

```
train on   log p_y  −  Σ_v q_v log p_v        instead of   log p_y
```

The gradient of the subtracted term is exactly `Σ_v q_v ∇log p_v`, so the expected update at every
position becomes zero: the drift is gone and no sample is thrown away. The engine only reports its
top-k log-probs per token, so the sum runs over those `k` tokens and the remaining probability mass
is approximated by a rescaled `p` (the paper's tail term; `k = 32` recovers ~95 % of the mass).

*With TIS or IcePop.* These keep the REINFORCE update but scale the sampled token by a per-token
weight `f(r)` of the train/rollout ratio `r = p_y / q_y` (TIS: `min(r, c)`; IcePop: `r` inside a
band, 0 outside). Score centering then subtracts the average of the *weighted* push,
`Σ_v q_v f(r_v) log p_v`. This stays exact because `f` only needs `(p, q)` at this position, so it
can be evaluated for every candidate `v`.

*FlashREINFORCE needs none of this.* Its recipe applies the exact ratio `r = p/q` to every surviving
token and drops whole responses whose train/rollout mismatch is too large (the `seq`/`geo` gate).
With the exact ratio the average push is `Σ_v q_v (p_v / q_v) ∇log p_v = ∇ Σ_v p_v = 0`: importance
sampling already removes the drift, and the gate only bounds the variance. Simple, and drift-free by
construction. For the same reason score centering cannot be stacked on it: with `f(r) = r` the
centering term is identically zero, and the gate decides keep-or-drop from later tokens, so no
per-candidate correction exists — `PolicyLoss` rejects the pair.

| You want | Flags |
|---|---|
| Score centering alone (replaces the IS correction) | `--algo.advantage.score_centering_top_k 32 --algo.advantage.is_correction_level off` |
| Score centering + TIS | `... --algo.advantage.is_correction_level token --algo.advantage.is_correction_mode trunc` |
| Score centering + IcePop | `... --algo.advantage.is_correction_level token --algo.advantage.is_correction_mode mask` |
| The `seq`/`geo` gate (FlashREINFORCE recipe) | keep the gate, do not add score centering (`PolicyLoss` rejects the pair) |

</details>

### ⚡ MTP rollout — speculative decoding

<details>
<summary>How the MTP draft is used, what it changes and which checkpoints support it</summary>

Checkpoints that ship a multi-token-prediction (MTP) head — e.g. **Qwen3.6-MoE**
(`mtp_num_hidden_layers: 1`) — can use it to **speed up generation** via vLLM
speculative decoding. One flag turns it on:

```bash
--vllm.mtp_num_speculative_tokens 1   # 0 = off (default); 1 is a good default
```

vLLM auto-detects the per-architecture MTP draft from the served checkpoint
(`qwen3_5_moe → qwen3_5_mtp`), drafts *N* tokens, and the target model verifies
each by rejection sampling. This is **lossless**: accepted tokens follow the
target policy's distribution, so the rollout log-probs stay unbiased for the RL
objective — it only changes throughput, never the learned policy.

Notes:
- **Rollout-only.** Like verl and NeMo-RL, the RL/SFT loss is **main-head-only**;
  Molt does not train the MTP head. The draft shares `embed_tokens`/`lm_head` with
  the target (refreshed every weight broadcast), so it tracks the updating policy;
  only the small MTP block stays at its checkpoint weights, so acceptance degrades
  gracefully rather than off a cliff.
- **Model support is vLLM-side.** Qwen3.6-MoE works out of the box. The official
  Nemotron-Nano-Omni HF checkpoint ships **no** MTP head, and vLLM 0.24 only
  auto-detects MTP for the *Super*-Omni arch (not Nano) — so omni3 rollout-MTP is
  unavailable until upstream adds it; vLLM errors at engine init if enabled on an
  unsupported checkpoint.

</details>

### 🎯 MoE routing stability — Router Replay (R3) and router freeze

<details>
<summary>Why MoE RL drifts, how R3 replays vLLM's routing, when to freeze the router</summary>

MoE RL is unstable because the rollout (vLLM) and training (FSDP) routers pick
experts **independently** — even at identical weights, numerical differences
flip a fraction of the top-k per layer, compounding until most tokens route to
different experts than they did during rollout. That breaks the importance-
sampling assumption behind GRPO/GSPO. Molt closes the gap at three levels — the
first two are on by default in the qwen3.5-moe recipes, the third is an optional
heavier alternative to R3:

**fp32 router precision (default).** The gate linear + expert-output combine run in
fp32 (matching vLLM's fp32 router) so the two sides agree on the gate *weights* to
begin with. A bf16 router silently drifts from vLLM and makes `vllm_kl` climb with
training. Override with `MOLT_GATE_PRECISION=bfloat16`.

**Rollout Routing Replay (R3, default)** — fix the top-k *selection* at the source
([arXiv:2510.11370](https://arxiv.org/abs/2510.11370)): vLLM returns the per-token
expert ids it chose, and the training forward replays that exact selection.

```bash
--train.routing_replay   # default in the qwen3.5-moe recipes
```

- **Freezes the routing, not the router.** Only the discrete top-k *selection* is
  replayed; the router logits are still recomputed from the live weights, so the
  gradient keeps flowing into the router (it keeps learning).
- **Full-sequence, absolute-position aligned**: routing is laid down by token
  position; any position the engine returns no routing for keeps its natural selection.
- Needs AutoModel `RouterReplay` (`nemo_automodel.components.moe.router_replay`,
  PR #2797). Incompatible with `--train.partial_rollout_enable` (vLLM frees routing
  on preemption).

**Router freeze (optional, blunter).** Hold the gate/router weights fixed so the
routing can't drift at all — excludes the router from the optimizer *and* the refit,
so vLLM and the actor route to the same experts **by construction**, no engine
support needed. The trade-off: the router stops learning, so it's **redundant with
R3** and off by default; reach for it only if routing drift still dominates and a
fixed router is acceptable.

```bash
--actor.freeze_moe_router   # off by default; redundant with R3
```

</details>

### 🧩 LoRA fine-tuning

<details>
<summary>Flags, target matching, learning rates and the RL constraints</summary>

Both paths take the same three flags (`--model.lora_*` for SFT, `--actor.lora_*` for RL);
`--*.lora_dim 0` (the default) is plain full fine-tuning:

```bash
--model.lora_dim 8 --model.lora_alpha 32 \
--model.lora_target_modules '*.q_proj' '*.k_proj' '*.v_proj' '*.o_proj'
```

Targets are wildcard matches against **full** module names (AutoModel's `ModuleMatcher`),
so a bare leaf name like `q_proj` silently matches nothing — omit the flag to patch every
linear layer instead. LoRA runs on bf16 master weights, so the full-fine-tune default LRs
are too small for it (SFT `5e-6`, RL `1e-6`): bf16 AdamW rounds those steps away and the
adapters barely move — molt warns at startup, use a LoRA-scale LR (`1e-4` and up). RL
trains the adapters only: the reference model keeps the plain checkpoint, which equals the
step-0 policy.

Constraints:
- **MoE experts and a tied `lm_head` are refused in RL.** The refit merges `scale·B@A`
  into the base weight before broadcasting, and neither AutoModel's grouped MoE adapters
  nor a tied `lm_head` (skipped in favour of `embed_tokens`) has such a mapping — the
  rollout would silently serve the untrained base.
- The merge lands on one bf16 weight rounding, so vLLM's weights differ from the
  trainer's by that bit. Harmless under IS correction; strictly on-policy runs should
  know the rollout policy is the trained policy only up to bf16 rounding.
- `--ckpt.save_hf` writes a **PEFT adapter directory** for LoRA runs
  (`adapter_model.safetensors` + `adapter_config.json`), not merged full weights; load it
  with `PeftModel.from_pretrained(base, adapter_dir)` or vLLM's `--enable-lora`. The DCP
  resume checkpoints stay full-weight, so resuming is unchanged.

</details>

## ✅ Validation

Fast local checks:

```bash
python -m compileall -q molt examples/python tests
pytest -q
```

Container checks:

```bash
SKIP_BUILD=1 DOCKER_GPUS=all DOCKER_SHM_SIZE=32g \
  bash examples/scripts/docker_run.sh "pytest -q"
```

## 🙏 Acknowledgement

Molt is based on OpenRLHF and keeps its Python package layout where
practical. The active architecture is intentionally minimal: a
Gymnasium-aligned agent (`Env` / `ChatAgent`), a single trainable actor,
optional KL reference workers, vLLM generation, and online policy
optimization — all on PyTorch + AutoModel.

## 📚 Citation

If you use Molt in your research, please cite:

```bibtex
@article{hu2026molt,
  title         = {Molt: A Scalable PyTorch-Native Training Framework for Agentic Reinforcement Learning},
  author        = {Jian Hu and Molt Contributors},
  year          = {2026},
  eprint        = {2607.21653},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG},
  url           = {https://arxiv.org/abs/2607.21653}
}
```

## 🤝 Contributing

External contributions are welcome — see [CONTRIBUTING.md](CONTRIBUTING.md).
All commits must be signed off (`git commit -s`) per the
[Developer Certificate of Origin (DCO)](https://developercertificate.org/).

## 📄 License

[Apache License 2.0](LICENSE). Copyright and third-party attributions:
[NOTICE](NOTICE) and [THIRD_PARTY_NOTICES.md](THIRD_PARTY_NOTICES.md).
