Skip to content

Release v0.1.0

Choose a tag to compare

View all tags
@hijkzzz hijkzzz released this 14 Jul 03:09
v0.1
0a64f71

Molt v0.1.0 — first public release 🦋

Molt is an agentic-first, PyTorch-native RL post-training framework: the agent is the program you write, the trainer is a single trainable actor, and the reward is arbitrary Python. Three components carry the whole runtime — Ray (placement + async queue), vLLM (rollout), and NVIDIA AutoModel + FSDP2 (training) — in ~9K lines of core RL code, no Megatron, that still scale to hundred-billion-parameter MoE.

📄 Technical report (DOI): 10.13140/RG.2.2.23375.65447 · 📖 DeepWiki · License: Apache-2.0

✨ Highlights

  • Agentic-first. Multi-turn tool use, multimodal (VLM) environments, and LLM-as-judge are first-class — the agent is a Gymnasium-style Env / ChatAgent, the reward is any Python you write.
  • PyTorch-native, no Megatron. Ray + vLLM + NVIDIA AutoModel/FSDP2. ~9K lines you can read end to end.
  • Scales. TP · EP · CP (· DP) — verified on Qwen3.5-397B, scaling toward GLM-5.2 753B, and DeepSeek-V3 at --fsdp.ep_size 256.
  • Train/inference consistency built in. First-class vllm_kl mismatch metric + importance-sampling correction, plus MoE router replay (R3) and router freeze.
  • Fully-async rollout with a vLLM router, partial rollout, and a one-flag switch back to strictly on-policy.

🧩 What's included

Training paradigms

  • SFT and the full RL stack, one CLI each (molt.cli.train_sft, molt.cli.train_rl_ray).
  • Estimators: reinforce, reinforce_baseline, rloo, grpo, dr_grpo, gae (PPO), on_policy_distill.
  • Optional KL reference workers (the reference doubles as the distillation teacher).

Agent contract

  • One Python module per run (--train.agent_path): choose Env (framework owns the LLM loop) or ChatAgent (you own it).
  • Rewards, graders, tools, VLM environments, and LLM-as-judge — all plain Python; judge calls route back through the same vLLM engines.

Scale & parallelism

  • FSDP2 with TP / EP / CP (· DP), comm/compute overlap, optimizer and full CPU offload.
  • 1T-class MoE on vLLM (DeepSeek-V3 at EP=256).

MoE routing stability

  • R3 (--train.routing_replay): replay vLLM's per-token expert selection in the training forward.
  • Router freeze (--actor.freeze_moe_router).

Off-policy / async correctness

  • vllm_kl metric for the rollout-vs-train logprob gap.
  • IS correction: --algo.advantage.is_correction_level {off,token,seq,geo} × --algo.advantage.is_correction_mode {mask,clip,trunc} (covers TIS, IcePop, seq-mask-tis).
  • Async rollout (--train.async_queue_size), partial rollout (--train.partial_rollout_enable).

Modalities & templates

  • Text and VLM (image) prompts, multi-turn.
  • Assistant-span / loss-mask derivation from the model's own chat template — verified on Qwen3.x, Nemotron Omni, Kimi, GLM, Gemma, and DeepSeek.

🚀 Getting started

Prebuilt image (full CUDA-13 stack — torch, vLLM, TransformerEngine, flash-attn, mamba, DeepEP, AutoModel):

docker pull hijkzzz/molt:latest   # or the pinned tag hijkzzz/molt:0.1

Minimal RL run:

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

Ready recipes from Qwen3-4B to 397B/753B-scale live under examples/scripts/.

📚 Links

📝 Notes

  • First public release; APIs may still change. Issues and PRs welcome (commits must be signed off, git commit -s, per the DCO).
  • Qwen3.5-397B is verified end to end; GLM-5.2 753B is the upper end of ongoing scaling.

Full details: see the README. Apache-2.0.

Assets 2
Loading