Molt: framework PyTorch para aprendizado por reforço agêntico
O que é o Molt
Molt é um framework de aprendizado por reforço (RL) agêntico desenvolvido pela NVIDIA, publicado sob o namespace NVIDIA-NeMo no GitHub. A proposta central é ser o menor stack PyTorch-nativo capaz de treinar modelos de até 1 trilhão de parâmetros (classe MoE, como DeepSeek-V3) com RL totalmente assíncrono, multimodal e multi-turno. O código de RL tem aproximadamente 9.200 linhas.
Arquitetura
O sistema se resume a três componentes: Ray para orquestração e filas assíncronas, vLLM para geração (rollout) e NVIDIA AutoModel + FSDP2 para treinamento em PyTorch puro. O contrato é "token-first": token ids, logprobs, faixas de ação, recompensas e tensores multimodais permanecem alinhados do rollout ao treinamento.
Design agentic-first
O agente é o programa. O pesquisador define a recompensa em Python puro dentro de uma classe Env (estilo Gymnasium, com step()/reset()) ou ChatAgent (que usa SDKs OpenAI ou Anthropic contra um servidor vLLM auto-lançado). O framework cuida de tokenização, orçamentos por turno e contabilidade multimodal. Qualquer coisa computável em Python é uma recompensa válida - incluindo LLM-as-judge.
Algoritmos e recursos
- Estimadores: REINFORCE, REINFORCE com baseline, RLOO, GRPO, DR-GRPO, GAE (PPO) e destilação on-policy.
- Correção de importance sampling (IS): compensa a divergência de logprobs entre rollout assíncrono e treinamento, com níveis (token, seq, geo) e modos (mask, clip, trunc). Cobre esquemas como TIS, IcePop e MIS.
- Estabilidade de roteamento MoE: Router Replay (R3) reproduz a seleção top-k do vLLM no forward de treinamento; router freeze congela os pesos do gate como alternativa mais bruta.
- MTP rollout: speculative decoding via cabeças multi-token prediction (ex.: Qwen3.6-MoE), lossless para o objetivo de RL.
- Destilação on-policy: um único flag transforma o modelo de referência em professor, usando KL reverso por token como sinal de treino.
Escala
Suporta tensor parallel, expert parallel (até EP 256), context parallel e offload de otimizador para CPU. O mesmo script que treina um modelo de 8B escala para MoE de 1T sem reescrita. Receitas prontas incluem Qwen3-4B, Qwen3.6-35B-A3B (VLM), Nemotron-Omni-30B e GLM-5.2 (~750B).
Comparação com alternativas
| Molt | OpenRLHF | verl | slime | |
|---|---|---|---|---|
| Backend | PyTorch/FSDP2 + AutoModel | DeepSpeed ZeRO-3 | FSDP/Megatron | Megatron |
| Código RL | ~9.2K LOC | ~7.2K | ~62K | ~25K |
| Foco | pesquisa agêntica | cobertura RLHF | amplitude produção | throughput Megatron |
Instalação
Disponível via Docker (imagem pré-construída com CUDA 13, torch 2.11, vLLM, TransformerEngine) ou PyPI (pip install molt-rl[vllm]). Licença Apache 2.0.