8月1日、Asif Razzaqが「NVIDIA AI Releases Molt: A PyTorch-Native Agentic Reinforcement Learning Framework」と題した記事を公開した。この記事では、NVIDIAのNeMoチームが開発したPyTorchネイティブのエージェント型強化学習フレームワーク「Molt」について詳しく紹介されている。
なぜ今NVIDIAがこのフレームワークを出すのか
LLMのポストトレーニング領域では、verl、OpenRLHF、slimeといったフレームワークが競合しており、近年急速に機能が拡充されている。その結果、コードベースは大規模化・複雑化し、「新しいアルゴリズムを素早く試す」という研究ニーズとの乖離が広がってきた。
NVIDIAはこの状況に対し、自社の強みであるH100/H200環境・FSDPとの深い統合を活かしつつ、「研究者が頭の中に保持できる」規模感を最優先にしたMoltを投入した。エージェント型の強化学習(マルチターンのツール使用・コード実行など)を想定した設計になっている点も、既存フレームワークとの差別化軸のひとつだ。
コードベースの小ささが最大の特徴
強化学習研究における実装コストは、研究者にとって長年の課題だ。新しいアルゴリズムを試すたびに、トレーナー、分散バックエンド、ロールアウトの接続部分を修正する必要があり、イテレーションのたびにその負担がのしかかる。
Moltはその課題に直接応える設計になっている。公式の説明によれば、「研究者が頭の中に保持でき、AIコーディングアシスタントがコードベース全体を読んで推論できる」サイズを目標としている。
実際の規模感を数字で見ると明確だ。各フレームワークのRLエントリポイントからインポートグラフをトレースして計測した結果は以下のとおりだ。
- Molt:約8,600行
- OpenRLHF:約7,200行(Moltより少ないが、後述する設計方針の違いから単純比較は難しい)
- slime:約25,000行
- verl:約62,000行
MoltとOpenRLHFはいずれもコンパクトな部類に入るが、設計方針は異なる。OpenRLHFがDeepSpeedやMegatronを含む幅広い環境をサポートする汎用的な構成を取るのに対し、MoltはFSDP2とvLLMに絞り込むことでシンプルさを維持している。verlとの比較では約7分の1というのがMoltの主要な訴求点だ。
ライセンスはApache 2.0。起動スクリプト、Slurmスクリプト、ビルド済みコンテナも同梱されている。
使える対象と用途
動作に必要なハードウェアは現実的に選ぶ必要がある。同梱のレシピは2ノード計16枚のH100 GPU(8枚をトレーニング、8枚をロールアウトに使用)を前提としている。
対象となるのは、フロンティアモデルを扱う研究所、ポストトレーニングを行うAIスタートアップ、金融・ヘルスケア・ロボティクス分野のエンタープライズAI研究グループ、マルチノードH100/H200アクセスを持つ学術ラボなどだ。
適用シーンとして記事が挙げているのは以下だ。
- マルチターンのツール使用エージェント
- コード実行エージェント
- 視覚言語環境(付属のgeo3kレシピ)
- LLM-as-judgeによる報酬ループ
- 小さなモデルへのオンポリシー蒸留
三つのコンポーネントを組み合わせた構成
Moltは既存のOSSをフォークせずに組み合わせる設計を採用している。具体的には、配置と非同期キューにRay、ロールアウトにvLLM、トレーニングにNVIDIA AutoModelとFSDP2を使う。いずれもフォークしていないため、上流の改善はコンテナのバージョンピンを更新するだけで取り込める。
ランタイムはエージェントプール、リクエストルーターを介したvLLMエンジン群、単一のトレーニング可能なポリシーアクターで構成される。ストリーミングプールがプロンプトグループを常に処理中の状態に保ち、アクターのトレーニング中もエンジンが止まらないようにする。部分ロールアウト時はエンジンを一時停止し、アクターのシャードをNCCL経由で各エンジンに直接ブロードキャストしてから、リクエストを破棄せず再開する。
エージェントはただのPythonプログラム
APIの設計も特徴的だ。RLの実行時にAgentRunnerをエクスポートするPythonモジュールを一つ指定するだけで、報酬関数も含め残りはすべて通常のコードで書ける。
二つの形式をサポートしている。Env形式ではGymnasium互換のstep()内でLLMループをフレームワークが管理する。ChatAgent形式ではOpenAIまたはAnthropicの標準SDKを通じてユーザー自身がループを制御できる。Moltはどちらのワイヤプロトコルにも対応したループバックサーバーを立ち上げ、すべてのリクエストをトークン単位で正確に集計する。
MoEにおける正確性の担保
設計上の三つの正確性不変条件のうち、**Mixture-of-Experts(MoE)モデル**に関わる「フォワード一貫性」が特に重要だ。ロールアウトとトレーニングのルーターがそれぞれ独立してエキスパートを選択すると、わずかな数値的差異がtop-kの選択を変えてしまう可能性がある。
Moltはこれに対し、ロールアウトルーティングリプレイを適用している。vLLMがトークンごとのエキスパートIDを返し、トレーニングのフォワードパスでそれを再現する仕組みだ。
スケール面では、4BパラメータのDenseモデルから700BのMoEまで、同じループを--fsdp.ep_size 256フラグ一つで動かせる。スループットはMegatronベースのスタックと統計的に同等とされているが、MoEのルーター不一致による注意点も開示されている。
詳細はNVIDIA AI Releases Molt: A PyTorch-Native Agentic Reinforcement Learning Frameworkを参照していただきたい。