NVIDIA-NeMo/labs-molt

What it solves

Moltは、高速なエージェンティックRL研究、特にフロンティア規模のMixture-of-Experts (MoE) モデル（最大1Tパラメータ）のトレーニングをターゲットに設計されています。小規模なモデルから大規模なモデルへとRLをスケールアップする際に通常伴う複雑さとインフラのオーバーヘッドに対処し、完全非同期のロールアウトとトレーニングをサポートしながら、ハック可能で読みやすいPyTorchネイティブなスタックを提供します。

How it works

Moltは、Rayを介した非同期ループによって調整される3つのコンポーネントからなるアーキテクチャを利用しています：

1. Ray: コンポーネント間の配置と非同期キューを管理します。
2. vLLM: ロールアウト（生成）フェーズを処理します。
3. NVIDIA AutoModel + FSDP2: 純粋なPyTorchでトレーニングフェーズを管理し、高度な並列化（TP, EP, CP）と大規模なアクターのためのAdam CPU offloadをサポートします。

報酬は、EnvまたはChatAgentクラス内のプレーンなPythonで定義されます。このフレームワークは「トークン優先」の契約を維持し、ロールアウトからトレーニングに至るまで、トークンID、logprobs、およびマルチモーダルテンソルが整合性を保つようにします。

Who it’s for

これは、エージェンティックな環境やRLアルゴリズムを大規模に、かつ迅速に反復させる必要があるAI研究者、特にVLM (Vision-Language Models) や大規模なMoEモデルを扱っており、8Bから1Tパラメータへのスケールアップ時にコードを書き直す必要がない人々に向けて構築されています。

Highlights

• Agentic-First Design: Gymnasiumに準拠したAPIを使用しており、エージェントがプログラムとなるため、研究者がトレーナーに触れることなくPythonで環境を反復させることができます。
• Frontier-Scale MoE: TP/EP/CPおよびMoEネイティブなトレーニングをネイティブにサポートし、DeepSeek-V3のような1Tクラスのモデルにスケールさせることが可能です。
• Fully-Async Runtime: ロールアウト、トレーニング、および重みの同期をオーバーラップさせることで、大規模なアクターにデータを供給し続けます。
• Small Codebase: RLコードが約9.2K行であり、読みやすく、エンドツーエンドでハックすることが容易です。
• Broad Algorithm Support: REINFORCE, RLOO, GRPO, PPO (via GAE), およびオンポリシーの蒸留を含みます。
• IS Correction: 非同期ロールアウトとトレーニング間のlogprobの不一致を処理するために、さまざまな重要度サンプリング補正スキーム（TIS, IcePop, MIS）を実装しています。