NVIDIA-NeMo/labs-molt

An agentic-first RL framework for research (9k lines).

解決的問題

Molt 是為高速度的代理型強化學習(RL)研究而設計,專注於訓練前沿規模的混合專家(MoE)模型(最高達 1T 個參數)。它解決了從小型模型擴展到大型模型時通常伴隨的複雜性與基礎設施開銷問題,提供一個原生 PyTorch 的堆疊,同時保持可修改性和可讀性,並支援完全非同步的推論與訓練。

如何運作

Molt 使用由 Ray 協調的三組件架構,透過非同步迴圈進行協調:

  1. Ray:管理組件之間的放置與非同步佇列。
  2. vLLM:處理推論(生成)階段。
  3. NVIDIA AutoModel + FSDP2:以純 PyTorch 管理訓練階段,支援先進的平行處理(TP、EP、CP)以及 Adam CPU 離線,以支援大型代理。

獎勵定義在 EnvChatAgent 類別中的純 Python 程式碼中。該框架維持「以 token 為先」的合約,確保從推論到訓練過程中,token ID、logprobs 和多模態張量始終對齊。

適用對象

專為需要在大規模下快速迭代代理環境與 RL 算法的 AI 研究人員設計,特別是那些使用 VLM(視覺-語言模型)和大型 MoE 模型的研究者,無需在從 8B 擴展到 1T 參數時重寫程式碼。

特色

  • 代理優先設計:採用與 Gymnasium 對齊的 API,其中代理即是程式,讓研究人員能在不觸碰訓練器的情況下,以 Python 快速迭代環境。
  • 前沿規模 MoE:原生支援 TP/EP/CP 和 MoE 原生訓練,可擴展至 DeepSeek-V3 類似的 1T 級模型。
  • 完全非同步執行時間:重疊推論、訓練與權重同步,以確保大型代理持續供應。
  • 小規模程式碼庫:約 9.2K 行 RL 程式碼,易於整體閱讀與修改。
  • 廣泛的演算法支援:包含 REINFORCE、RLOO、GRPO、PPO(透過 GAE),以及在策略蒸餾。
  • IS 修正:實作多種重要性採樣修正方案(TIS、IcePop、MIS),以處理非同步推論與訓練之間的 logprob 不匹配問題。

相關

  • 專案
  • 專案
  • 專案
  • 專案
  • 專案