NVIDIA-NeMo/labs-molt

Molt 是**面向 Agentic 优先**且**原生支持 PyTorch** 的。Agent 即程序； 训练器是一个单一的 actor；奖励是你在 `Env` 或 `ChatAgent` 中编写的任何 Python 代码 —— 包括评分器、多轮工具、VLM 环境和 LLM-as-judge。 其余部分由三个组件承载 —— **Ray** 负责部署和异步队列， **vLLM** 负责 rollout，**NVIDIA AutoModel + FSDP2** 负责在纯 PyTorch 中进行训练。 这就是整个技术栈：**约 9.2K 行 RL 代码**，可在带有 TP / EP / CP 的 vLLM 上 扩展到 1T 级别的 MoE** —— 例如在 `--fsdp.ep_size 256` 下运行 DeepSeek-V3， 并为最大的 actor 使用 Adam CPU offload。一个 Agent API，一个可训练的 actor，简洁到足以让你从头到尾通读。 ## 🧩 架构 三个模块。一个异步循环。

**Ray** 掌控着这三个模块之间的部署和异步队列 —— 这就是 整个运行时。其契约是**以 token 为先**：token id、 logprob、动作范围、奖励和多模态 tensor 从 rollout 到训练始终保持对齐。任何你能在 Python 中计算的内容都可以作为奖励， 包括通过驱动 rollout 的同一批 vLLM 引擎进行 LLM-as-judge 调用。 ## ✨ 为什么选择 Molt | | 你能得到什么 | 为什么这对研究很重要 | |---|---|---| | 🤖 **Agentic 优先** | 一个与 Gymnasium 对齐的 API —— `Env.step()` 或 `ChatAgent.run()` —— 涵盖评分器、多轮工具、VLM 环境以及兼容 OpenAI/Anthropic 的服务器 | Agent *即*程序 —— 用纯 Python 迭代环境，而训练器保持不变 | | ⚙️ **全异步运行时** | Ray 部署、异步 rollout 队列、vLLM 引擎、部分 rollout、权重同步 | Rollout、训练和权重同步重叠进行 —— 一个 DeepSeek-V3 级别的 actor 无需定制基础设施即可保持数据供给 | | 🔥 **PyTorch 原生，AutoModel 优先** | FSDP2 + NVIDIA AutoModel，端到端纯 PyTorch | 用你已经在写的语言修改模型；没有后端的繁文缛节 | | 🎯 **单 actor 的极简性** | 一个 actor，可选的 KL 参考模型 —— 整个 RL 计算图能画在一页纸上 | 每个梯度都是显式的；每一个 loss 项都近在一个文件之内 | | 🚀 **前沿规模的 MoE** | AutoModel + FSDP2 + TP / EP / CP + Adam CPU offload，原生支持 MoE —— 例如在 `--fsdp.ep_size 256` 下运行 DeepSeek-V3 | 训练 8B 模型的同一个脚本可以扩展到 1T 级别的 MoE —— 规模切换无需重写代码 | | 🔗 **以 token 为先的契约** | 对齐的 token id、logprob、动作范围、奖励和多模态 tensor | 多轮、VLM 和工具调用轨迹端到端共享一种格式 | | 🪶 **小巧、易于修改的表层** | 跨 3 个薄层的约 9.2K 行 RL 代码 | 修改一层而无需触及其他层 —— 一个下午就能读完 | ## 📊 对比分析 RL 生态系统为广度而优化。Molt 则为**大规模 Agentic 研究的速度**而优化 —— 这是最小的 PyTorch 原生技术栈，仍能在 vLLM 上以前沿的 MoE 规模驱动全异步的 Agentic RL。 | | **🦋 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 | **一个框架，一项任务。** Molt 是最小的 PyTorch 原生 技术栈，能够在 vLLM 上将 NVIDIA AutoModel 从 SFT 推进到前沿规模的 Agentic RL。用纯 PyTorch 阅读每一行触及你梯度的代码。 ## 🎯 支持范围 ### 训练与运行时 | 领域 | 支持 | |---|---| | SFT | `molt.cli.train_sft` | | RL | 通过 `molt.cli.train_rl_ray` 实现由 vLLM 支持的在线 RL | | Runtime | Ray 部署、异步 rollout 队列、vLLM 引擎、部分 rollout 同步 | | Model scale | AutoModel + FSDP2，带有 TP / EP / CP，原生支持 MoE —— 例如在 `--fsdp.ep_size 256` 下的 DeepSeek-V3 | | Model backend | **NVIDIA AutoModel 是主要路径** —— 原生 CP / EP / TP，自定义 MoE+EP 并行器，TE 融合注意力；模型端的一切都与 AutoModel 自身的配方保持一致。HF transformers 路径是**非首选的备选方案**（仅当模型没有原生类时，AutoModel 才会降级到该方案），仅支持 **text + flash_attention_2 + packing —— 不支持 CP / EP / TP** | | Optimizer | `adam`（默认），对于最大的 actor 提供 CPU offload (`--fsdp.offload optimizer`)。`muon`（通过 Dion 实现 Newton–Schulz：对 2D 权重和分组的 MoE 专家使用 Muon，对 embeddings / head / norms 使用 AdamW）目前为**实验性** —— 可以分布式运行 (FSDP / EP)，但尚未表现出比 `adam` 持续更好的效果，因此 `adam` 仍然是推荐的默认选项 | ### Agent 与奖励 | 领域 | 支持 | |---|---| | Agent interface | 通过 `--train.agent_path` 指定 `Env` 或 `ChatAgent` 子类 + 一个 `AgentRunner` | | Reward source | 从 `Env.step` 或 `ChatAgent.run` 返回的 `Result(reward=...)` | | Modalities | 文本和 VLM 提示词，包括图像 payload | | Chat templates | Assistant 片段（SFT loss 掩码 + 多轮 rollout 拼接）衍生自模型自身的 chat template —— 没有硬编码的标记。已在 ChatML (Qwen3.x, Nemotron omni3)、Kimi-K2.6、GLM、Gemma 和 DeepSeek 上完成验证 | ### 算法 | 领域 | 支持 | |---|---| | Estimators | `reinforce`, `reinforce_baseline`, `rloo`, `grpo`, `dr_grpo`, `gae` (PPO), `on_policy_distill` | | PPO critic | `--algo.advantage.estimator gae` 会添加一个价值模型：它拥有自己的 Ray 组 (`CriticModelActor`)，默认情况下与 actor 的 GPU 协同部署，或者可解耦部署，提供 GAE 优势 (`--algo.advantage.lam`) + 裁剪的价值 loss (`--critic.value_clip`)，拥有自己的 optimizer/LR (`--critic.adam.lr`) 以及可恢复的 `_critic` 检查点。它基于 `NeMoAutoModelForCausalLM` + 一个标量价值头构建，因此保留了原生的 TP / EP / CP 路径 | | Distillation | On-policy 蒸馏 —— 针对冻结的 teacher 模型计算逐 token 的反向 KL，通过 `--algo.advantage.estimator on_policy_distill` + `--ref.model_name_or_path` 实现 | | IS correction | 针对离线策略 / 异步 rollout 的训练/rollout logprob 不匹配进行重要性采样校正：`is_correction_level {off,token,seq,geo}` × `is_correction_mode {mask,clip,trunc}`（涵盖 TIS、IcePop、seq-mask-tis；参见下文的 *IS correction*） | | KL | 当 `--algo.kl.init_coef > 0` 时启用可选的 reference worker（reference 模型同时兼任蒸馏的 teacher） | ### MoE 路由稳定性 | 领域 | 支持 | |---|---| | Router replay (R3) | `--train.routing_replay` —— vLLM 中逐 token 的 top-k 选择在训练前向传播中被重放；详情见 Scaling Knobs 下的 *MoE 路由稳定性*部分 | | Router freeze | `--actor.freeze_moe_router` 固定 gate/router 权重，使得 vLLM 和 actor 将 token 路由到相同的专家。这稳定了 MoE RL / 蒸馏，并缩小了 IS 校正过滤器所针对的 rollout 与训练之间的 logprob 差距 —— 在每次重新拟合之间发生漂移的 router 是造成该差距的主要原因 | ### IS 校正（训练/rollout 的 logprob 不匹配） 异步和部分 rollout 会导致 FSDP actor 重新计算的 `pi_train` 偏离 vLLM 生成时的 `pi_rollout`（不同的 kernel，加上 HTTP router 无法察觉的请求中期的权重交换）。Molt 使用 逐 token 的重要性比率 `pi_train / pi_rollout` 校正由此产生的离线策略更新，该比率受两个旋钮控制： - `--algo.advantage.is_correction_level {off, token, seq, geo}` —— 受限比率的粒度。`off` 禁用校正；`token` 限制每个 token 自身的比率； `seq`/`geo` 将一个序列的比率聚合 (`exp(sum)` / `exp(mean)`)一个序列级别的 **拒绝过滤器**（保留的序列仍然带有其逐 token 的 IS 权重），因此它们需要 `mode mask`。 - `--algo.advantage.is_correction_mode {mask, clip, trunc}` —— 对超出范围的单位 (unit) 的处理方式。`mask` 将其丢弃（零梯度）；`clip` 将其权重钳制在范围内；`trunc` 仅截断上尾。 - `--algo.advantage.is_correction_threshold LOW HIGH` —— 比率上的 `[low, high]` 范围（配方使用紧凑的 `0.99 1.01`）。 命名方案及其现有技术： | 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 | 参考文献：**TIS**（训练/推理比率的截断重要性采样），**IcePop** （超出范围比率的 token 级别掩码），以及 **MIS**（掩码重要性采样，Yingru Li —— 序列级别的掩码 IS，这启发了 `seq`/`geo` 拒绝过滤器）。 ## 📦 安装 首先克隆此代码库 —— 启动脚本、agent 和配方都在这里， 并且 `examples/scripts/docker_run.sh` 会将此检出挂载到容器中。对于本地 （非容器）开发，请添加可编辑安装：它会拉取本仓库验证过的精确锁定 git 版本的 AutoModel，因此 R3 路由重放和 Muon 开箱即用： ``` 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 ``` **推荐的路径是使用项目容器**（`dockerfile/Dockerfile`）。它打包了完整的 CUDA-13 技术栈 —— torch 2.11 · vLLM · TransformerEngine · flash-attn · mamba · DeepEP · NVIDIA AutoModel —— 专为 A100 / H100 / H200 / B200·GB200 构建，因此它可以直接运行 SFT 和 RL，而无需在本地费力处理依赖关系。从 Docker Hub 拉取预构建的镜像： ``` docker pull hijkzzz/molt:latest # or a pinned release: hijkzzz/molt:0.1.3 ``` ...或者从 Dockerfile 自行构建（例如为了更改 CUDA / vLLM / AutoModel 的版本锁定）： ``` docker build -f dockerfile/Dockerfile -t hijkzzz/molt:latest . ``` 发布的包也位于 PyPI 上，支持免检出安装： ``` pip install "molt-rl[vllm]" ``` ## 🚀 快速开始 ### SFT ``` 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 使用与 RL 相同的 AutoModel/FSDP2 模型加载路径。 ### RL ``` 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 ``` 常见的 RL 开关： | 目标 | Flags | |---|---| | 禁用 reference worker | `--algo.kl.init_coef 0` | | 启用 KL 正则化 | 将 `--algo.kl.init_coef` 设置为大于零，并通过 `--ref.num_nodes`、`--ref.num_gpus_per_node` 或 `--train.colocate_fsdp_models` 部署 reference worker | | 每个提示词比较样本 | `--rollout.n_samples_per_prompt 8` 加上 `reinforce_baseline`、`rloo`、`grpo` 或 `dr_grpo` | | 解耦 rollout 和训练 | `--train.async_queue_size 2` | | 在同步期间保持 rollout 活跃 | `--train.partial_rollout_enable` | | 根据 agent 分数进行过滤 | `--algo.dynamic_filtering_enable --algo.dynamic_filtering_range 0.0 1.0` | | 校正异步 rollout logprob | `--algo.advantage.is_correction_level geo` (seq-mask-tis；token 级别需添加 `--algo.advantage.is_correction_mode clip/trunc/mask`) | | 冻结 MoE 路由（稳定 MoE RL） | `--actor.freeze_moe_router` | | On-policy 蒸馏 | `--algo.advantage.estimator on_policy_distill --ref.model_name_or_path /path/to/teacher` | | 独立的评估采样 | `--eval.temperature`, `--eval.top_p`, `--eval.max_new_tokens`, `--eval.n_samples_per_prompt` （未设置的将回退到 rollout） | | 评估检查点（不进行训练） | `--eval.eval_only --eval.dataset ` —— 对评估集进行一次评分然后退出；vLLM 持有 HF 权重，因此 policy/ref/critic 的 FSDP actor 永远不会被构建，它们的 GPU 将全部用于评估 | | 导出 / 重放 rollout 批次 | `--train.rollout_dump_dir `，然后使用 `--train.rollout_replay_dir ` 在不重新生成的情况下重新运行训练 | | 检查权重更新覆盖率 | `--train.check_weight_update_equal` 会警告哪些 vLLM 参数在广播后处于陈旧状态 | ## 🤖 Agent 契约 每次 RL 运行都会指向一个 Python 模块： ``` --train.agent_path /path/to/agent.py ``` 该模块必须导出 `AgentRunner`。在两条路径中**选择其一**： ### 1. `Env` — 框架掌管 LLM 循环 *（Gymnasium 风格的 step/reset）* ``` 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) ``` 框架驱动 vLLM、分词、多模态核算以及 每轮的预算。你的 `step()` 返回一个 `Result`；框架会连续执行轮次，直到 `terminated` 或 `truncated`。 ### 2. `ChatAgent` — 你通过 OpenAI **或** Anthropic SDK 掌管循环 ``` 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) ``` 同一个服务器也能解析 Anthropic 协议 —— 将 `AsyncAnthropic` 指向 `ctx.session_url`（session 根路径*不包含* `/v1`；SDK 会自动追加 `/v1/messages`），其他一切都完全相同： ``` 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 ``` 一个 FastAPI vLLM 服务器会在 session URL 下的环回接口自动启动。 外部 HTTP 调用者（浏览器自动化、评估测试框架、OSWorld 等）通过 `/v1/chat/completions` (OpenAI) 或 `/v1/messages` (Anthropic) 访问同一个引擎 —— 两者都会解码为严格的逐 token 累加。 #### Context 压缩 → 每次 rollout 生成多个 step-sample 每一次聊天调用都会将前一轮**精确的** token 向前传递，并且只追加 新的增量 (delta)，因此一个多轮的 episode 会拼接成一条单调递增且 token 精确的轨迹。但是长周期的 agent 通常会**压缩**其上下文 —— 总结 或丢弃旧的轮次以保持在窗口限制内（例如 `/compact` 步骤）—— 这会*重写*前缀，因此它不再是已分词内容的干净扩展。 模型自身的 chat template 也可以重写前缀：一旦后续有新的 user 查询，Qwen3 风格的模板就会重新渲染先前的 assistant 轮次，且不带 `` 块。 服务器会自动检测到这一点：当传入的请求重写了 前缀而不是对其进行扩展时，它会**封存当前片段并从重新模板化后的对话中开启一个全新的、token 精确的片段**。因此，一次 rollout 会生成多条片段轨迹 —— 它们共享这次 rollout 的 奖励和 `rollout_id`，因此组别基线 (GRPO/RLOO/…) 会将它们去重为*每次 rollout 一个奖励*，同时每个片段仍然会将自身生成的 token 贡献给给策略梯度（这与多轮 agent 使用的 step-sample 契约相同）。无需 在 agent 端进行任何更改 —— 它在两种协议上均适用，包括对于我们来说压缩过程是不透明的外部框架 （Claude Code, opencode, AgentScope 等）。 ### `Result` 字段 | 字段 | 含义 | |---|---| | `reward` | 训练器消耗的标量奖励（必填） | | `observation` | 下一轮的观察文本（仅多轮场景） | | `terminated` | Episode 自然结束；默认值为 `True` | | `truncated` | 被外部截断（最大轮次、长度等） | | `info` | 用于记录的可选标量诊断字典 | | `score` | 可选的动态过滤 / 仪表板分数（默认为 `reward`） | | `images` | 可选的下一轮图像列表 | | `sampling_params` | 可选的每轮覆盖参数 | 在 `examples/python/agents/` 下提供了四个参考 agent： ``` --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 ``` ## 🍳 配方 参考启动脚本位于 `examples/scripts/` 下。目前提供两个端到端的系列，均基于 AutoModel + FSDP2 后端： | Workflow | quick_start | slurm | |---|---|---| | 在 geo3k 上进行 Qwen3.6-35B-A3B VLM SFT | `quick_start/sft_qwen3_6_35b.sh` | `slurm/sft_qwen3_6_35b.sh` | | 在 geo3k 上进行 Qwen3.6-35B-A3B VLM RL (多轮 Python 工具) | `quick_start/rl_qwen3_6_35b.sh` | `slurm/rl_qwen3_6_35b.sh` | | 在文本数学任务上进行 Qwen3-4B dense SFT | `quick_start/sft_qwen3_4b.sh` | `slurm/sft_qwen3_4b.sh` | | 在文本数学任务上进行 Qwen3-4B dense RL | `quick_start/rl_qwen3_4b.sh` | `slurm/rl_qwen3_4b.sh` | | 在 geo3k 上进行 Nemotron-Omni-30B-A3B VLM RL (混合 SSM MoE, CP8+EP8) | — | `slurm/rl_omni3_30b.sh` | | Nemotron-Omni-30B-A3B on-policy 蒸馏 | — | `slurm/rl_distill_omni3_30b.sh` | | 在文本数学任务上进行 GLM-5.2 ~750B RL (MLA + DSA sparse attention, EP256) | — | `slurm/rl_glm5_2_753b.sh` | 单节点快速入门用法： ``` MODEL_PATH=/path/to/Qwen3-4B bash examples/scripts/quick_start/rl_qwen3_4b.sh ``` geo3k VLM 脚本（`rl_qwen3_6_35b.sh` / `sft_qwen3_6_35b.sh`）会在首次运行时通过 `examples/python/utils/prepare_geo3k.py` 自动准备 数据集。要手动预暂存它 （或刷新它），请运行： ``` python3 examples/python/utils/prepare_geo3k.py --num-proc 8 --out-dir .tmp/geo3k ``` 或者将 `PROMPT_DATASET` / `EVAL_DATASET` 指向你自己的数据。 Slurm 用法： ``` # 1) 在 2 个 interactive 节点上进行 SFT smoke 测试 sbatch examples/scripts/slurm/sft_qwen3_6_35b.sh # 2) 在 2 个 interactive 节点上进行 RL smoke 测试（首次运行时自动准备 geo3k） sbatch examples/scripts/slurm/rl_qwen3_6_35b.sh # 3) 将 RL 扩展到 4 个节点以实现 convergence sbatch --nodes=4 examples/scripts/slurm/rl_qwen3_6_35b.sh ``` ### 多轮 Python 工具环境 `examples/python/agents/geo3k.py` 是由 Qwen3.6 RL 脚本使用的 VLM 多轮配方。模型发出调用 `python_executor(code=...)` 的 ``；环境会在沙盒子进程中运行该代码片段，并将捕获到的 stdout 作为 `` 轮次反馈回去。 该循环最多运行 `MAX_AGENT_TURNS` 次（agent 默认为 5；附带的 Qwen3.6 配方将其提升至 10）；最终的 `ANSWER`（或针对旧版本的 `\boxed{ANSWER}`）会与标准答案进行评分对比，并作为奖励。 ### 兼容 OpenAI / Anthropic 的服务器 Agent 对于已经能够处理 OpenAI Chat Completions 或 Anthropic Messages API 的 agent，请继承 `ChatAgent` 子类（见 `examples/python/agents/chat_minimal.py`）。自动启动的服务器针对滚动更新的 vLLM 引擎同时暴露了 `/v1/chat/completions` 和 `/v1/messages`，因此任何外部循环（浏览器自动化、评估 测试框架、OSWorld 等）都可以通过原生的 OpenAI 或 Anthropic SDK 驱动策略 —— 两种协议都会解码为同样的、精确到 token 的轨迹捕获。 ### On-policy 蒸馏 在 student 模型*自身*的 on-policy 样本上，将其蒸馏到冻结的 teacher 模型上。只需一个开关 —— `--algo.advantage.estimator on_policy_distill` —— 即可将 reference 模型变成 teacher 模型，并将针对它的逐 token **反向 KL** 作为整个训练 信号：优势函数变为 `-kl_coef · (log π_student − log π_teacher)`，没有标量奖励，没有组别基线，也没有白化，因此策略 loss 就是反向 KL 梯度的策略梯度估计量，它会将 student 模型拉向 teacher 模型。 ``` 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 ``` 其他一切都源自这一个开关，因此纯蒸馏不需要 奖励函数，也不需要任务 agent —— 只需要 teacher 检查点。选择该估计器会强制使用 `--algo.kl.estimator k1`，关闭 `--algo.kl_loss`（KL 通过优势函数流动，而不是作为一个单独的 loss 项），将 `--algo.kl.init_coef` 默认为 `1.0`，并且 —— 当未提供 `--train.agent_path` 时 —— 会自动选择一个内置的单轮、兼容 VLM 的生成器 (`molt/agents/distill_agent.py`)，该生成器为每个提示词采样一个 on-policy 的补全，并返回一个估计器会忽略的虚拟 `0.0` 奖励。 Teacher 模型**必须共享 student 模型的 processor/tokenizer**，以便逐 token 的 logprob 能在相同的（视觉展开后的）序列上对齐 —— 通常是来自同一系列更大或训练更充分的检查点。它以仅推理的方式加载，可以 与 actor 节点协同部署 (`--train.colocate_fsdp_models`)，或者分配其自己的 `--ref.num_nodes`。观察 `kl` / `logprobs_diff` 随着 student 模型匹配 teacher 模型而下降至 0；任务准确率不是目标，因此评估是关闭的。 要蒸馏**多轮工具使用分布**（与 student 模型实际部署的方式相匹配），请将 `--train.agent_path` 指向任务的 真实 agent（例如 `chat_geo3k.py`）—— 其奖励只会被估计器忽略。 `examples/scripts/slurm/rl_distill_omni3_30b.sh` 是基于 omni3 EP8 / CP8 / TE / DeepEP 配方的现成 VLM 示例。 ## 🎛️ 扩展配置旋钮 Molt 的目标是结合 FSDP2 的 AutoModel 自定义模型： | | 模式 | 标志 | |---|---|---| | **Actor** | Tensor parallel | `--fsdp.tp_size 2` | | | Expert parallel | `--fsdp.ep_size 8` (例如对于 DeepSeek-V3 级别的 MoE 使用 `256`) | | | Context parallel | `--fsdp.cp_size 8` (32K+ 序列)，包括 VLM 和 MoE 路由重放 | | | Optimizer CPU offload | `--fsdp.offload optimizer` (为最大的 actor 释放 VRAM) | | **vLLM rollout** | Tensor parallel | `--vllm.tensor_parallel_size 2` | | | Expert parallel | `--vllm.enable_expert_parallel` (EP = TP × DP) | | | Data parallel | `--vllm.data_parallel_size 4` (单节点 mp；在 TP 之上提升 EP —— DeepSeek-V3 风格的 TP8+DP4 → EP32) | | | Scheduler token budget | `--vllm.max_num_batched_tokens 32768` | | | MTP spec-decode | `--vllm.mtp_num_speculative_tokens 1` | | **MoE 稳定性** | Router replay (R3) | `--train.routing_replay` | | | Router freeze | `--actor.freeze_moe_router` | Context 并行被委托给 AutoModel 的 `ContextParallelSharder`，因此每个 模型都能获得其注意力后端所需的分片 —— 混合 SSM / 线性注意力模型（Nemotron Omni, Qwen3.5-MoE）使用轮询调度，而稀疏 注意力模型 (GLM-5.2 DSA) 则使用扁平的 THD 流。VLM 视觉塔和路由重放会随着序列进行分片，因此 `--fsdp.cp_size` 可以与 `--data.image_key` 和 `--train.routing_replay` 组合使用。样本打包 (`--fsdp.packing_samples`) 仅限文本且默认关闭；在 CP 下它采用 THD 路径。 ### ⚡ MTP rollout（推测解码） 附带多 token 预测 (MTP) 头的检查点 —— 例如 **Qwen3.6-MoE** (`mtp_num_hidden_layers: 1`) —— 可以利用它通过 vLLM 推测解码来**加速生成**。只需一个标志即可开启： ``` --vllm.mtp_num_speculative_tokens 1 # 0 = off (default); 1 is a good default ``` vLLM 会自动从提供的检查点中检测到基于特定架构的 MTP 草稿 (`qwen3_5_moe → qwen3_5_mtp`)，生成 *N* 个 token，然后目标模型通过拒绝采样对每一个进行验证。这是 **无损的**：被接受的 token 遵循目标策略的分布，因此对于 RL 目标来说，rollout 的 logprob 保持无偏 —— 它只会改变吞吐量，永远不会改变学习到的策略。 注意： - **仅限 rollout。** 与 verl 和 NeMo-RL 一样，RL/SFT 的 loss **仅限主头**； Molt 不会训练 MTP 头。草稿模型与目标模型共享 `embed_tokens`/`lm_head`（在每次权重广播时都会刷新），因此它会跟踪不断更新的策略； 只有小型的 MTP 块保留其检查点权重，因此接受率会优雅地下降，而不是发生断崖式下跌。 - **模型支持取决于 vLLM 端。** Qwen3.6-MoE 开箱即用。官方的 Nemotron-Nano-Omni HF 检查点附带 **没有** MTP 头，并且 vLLM 0.24 仅能自动检测 *Super*-Omni 架构（而不是 Nano）的 MTP —— 因此在上游添加此功能之前，omni3 的 rollout-MTP 是 不可用的；如果在不受支持的检查点上启用，vLLM 会在引擎初始化时报错。 ### 🎯 MoE 路由稳定性 — 路由重放 (R3) 与路由器冻结 MoE RL 是不稳定的，因为 rollout (vLLM) 和训练 (FSDP) 的路由器是**独立**选择 专家的 —— 即使在权重完全相同的情况下，数值差异 也会导致每层一小部分的 top-k 发生翻转，不断累积，直到大多数 token 被路由到与 rollout 期间不同的专家。这破坏了 GRPO/GSPO 背后的重要性 采样假设。Molt 在三个层面上弥补了这一差距 —— 前两个在 qwen3.5-moe 配方中默认开启，第三个是 R3 的一个较重的可选替代方案： **fp32 router 精度（默认）。** Gate 线性层和专家输出结合在 fp32 中运行（与 vLLM 的 fp32 路由器相匹配），因此双方从一开始就在 gate *权重*上达成一致。bf16 路由器会悄悄偏离 vLLM，并导致 `vllm_kl` 随着训练不断攀升。可以使用 `MOLT_GATE_PRECISION=bfloat16` 进行覆盖。 **Rollout 路由重放 (R3，默认)** —— 从源头修复 top-k 的*选择* ([arXiv:2510.11370](https://arxiv.org/abs/2510.11370))：vLLM 返回它所选择的逐 token 的 专家 ID，训练的前向传播重放该确切的选择。 ``` --train.routing_replay # default in the qwen3.5-moe recipes ``` - **冻结的是路由，而不是路由器。** 只有离散的 top-k *选择*被 重放；路由器 logits 仍然是从实时权重中重新计算的，因此 梯度会继续流入路由器（它会继续学习）。 - **全序列，绝对位置对齐**：路由是由 token 位置决定的；引擎没有为其返回路由的位置会保留其自然选择。 - 需要 AutoModel 的 `RouterReplay` (`nemo_automodel.components.moe.router_replay`, PR #2797)。与 `--train.partial_rollout_enable` 不兼容（vLLM 在抢占时释放路由信息）。 **路由器冻结（可选，较为生硬）。** 将 gate/router 权重固定，使得路由完全无法漂移 —— 将路由器从优化器和重新拟合中排除，因此 vLLM 和 actor **通过构造**就能将 token 路由到相同的专家，不需要引擎的支持。权衡之处在于：路由器停止学习，因此它对于 R3 来说是**冗余的**，并且默认关闭；只有在路由漂移仍然占据主导，并且固定路由器是可以接受的情况下才使用它。 ``` --actor.freeze_moe_router # off by default; redundant with R3 ``` ## ✅ 验证 快速的本地检查： ``` python -m compileall -q molt examples/python tests pytest -q ``` 容器检查： ``` SKIP_BUILD=1 DOCKER_GPUS=all DOCKER_SHM_SIZE=32g \ bash examples/scripts/docker_run.sh "pytest -q" ``` ## 🙏 致谢 Molt 基于 OpenRLHF，并在可行的情况下保留了其 Python 包布局。活跃的架构是有意保持极简的：一个 与 Gymnasium 对齐的 agent (`Env` / `ChatAgent`)，一个可训练的 actor， 可选的 KL reference worker，vLLM 生成，以及在线策略 优化 —— 全部基于 PyTorch + AutoModel。 ## 📚 引用 如果你在研究中使用了 Molt，请引用： ``` @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.md](CONTRIBUTING.md)。 所有的 commit 必须被签名 (`git commit -s`)，这符合 [Developer Certificate of Origin (DCO)](https://developercertificate.org/)。 ## 📄 许可证 [Apache License 2.0](LICENSE)。版权及第三方署名： [NOTICE](NOTICE) 和 [THIRD_PARTY_NOTICES.md](THIRD_PARTY_NOTICES.md)。

标签：AI智能体, PyTorch, vLLM, 人工智能, 凭据扫描, 大模型训练, 强化学习, 用户模式Hook绕过, 逆向工具