AI News HubLIVE
サイト内リライト3 分で読了

NVIDIA AI、PyTorchネイティブのエージェント強化学習フレームワーク「Molt」をリリース

NVIDIAのNeMoチームは、約8.6K行のRLコードからなるPyTorchネイティブのエージェント強化学習フレームワーク「Molt」を公開した。Ray、vLLM、NVIDIA AutoModelを1つの非同期ループに組み合わせ、エージェントを通常のPythonコードとして扱い、トークン正確な軌跡を維持し、Megatronベースのスタックと統計的に同等のスループットを実現する。

ソースMarkTechPost著者: Asif Razzaq

エージェント強化学習の研究では、アルゴリズムの修正が絶え間なく行われます。新しい推定器、新しいパイプライン段階、新しいロールアウト方式が次々に登場します。主流のフレームワークでは、変更のたびにトレーナー、分散バックエンド、ロールアウトの接着層をすべて貫通する必要があり、研究者は反復のたびにそのコストを負担します。

NVIDIAのNeMoチームが公開したMoltは、このコストに直接対処します。これはPyTorchネイティブのエージェントRLフレームワークであり、研究者が全体を頭の中で把握でき、AIコーディングアシスタントが全体を読んで推論できるほどコードベースをコンパクトにするという設計目標を掲げています。RLコードの量は約8.6K行(各フレームワークのRLエントリポイントからのインポートグラフを追跡して測定)で、比較としてverlは約62K行、slimeは約25K行、OpenRLHFは約7.2K行です。

Moltはデプロイ可能でしょうか。可能です。Apache 2.0ライセンスで提供され、起動スクリプト、Slurmスクリプト、事前構築済みコンテナが同梱されています。ただし、研究論文では本番のトレーニングサービスではなく研究インフラと位置付けられており、実際の制約はハードウェアです。標準のレシピは8基のH100 GPUを搭載した2ノードを想定し、8基をトレーニング、8基をロールアウトに割り当てます。

これにより、Moltはフロンティアおよびそれに準ずる研究所、ポストトレーニングを行う資金力のあるAIスタートアップ、独自環境でエージェントを訓練する金融・医療・ロボティクスの企業研究グループ、マルチノードH100/H200を持つ学術研究所が利用できます。用途は、マルチターンのツール使用エージェント、コード実行エージェント、視覚言語環境(同梱のgeo3kレシピ)、LLM-as-judgeの報酬ループ、小型の学生モデルへのオンポリシー蒸留などです。

MoltはRay(配置と非同期キュー)、vLLM(ロールアウト)、NVIDIA AutoModel(FSDP2トレーニング)を組み合わせています。いずれもフォークされておらず、アップストリームの改善はリベースではなくコンテナピンとして取り込まれます。ランタイムはエージェントプール、リクエストルーターの背後にある一連のvLLMエンジン、そして単一の学習可能なポリシーアクターで構成されます。ストリーミングプールはプロンプトグループを飛行中に維持し、アクターのトレーニング中にエンジンが空にならないようにします。部分的なロールアウトではエンジンを一時停止し、NCCL経由でアクターのシャードを各エンジンに直接ブロードキャストしてから、破棄せずに保持したリクエストを再開します。

エージェントは通常のプログラムです。RL実行では、AgentRunnerをエクスポートする1つのPythonモジュールを指定するだけで、報酬を含めてすべてが通常のコードです。2つの形式がサポートされています。Envでは、フレームワークがGymnasium準拠のstep()でLLMループを所有します。ChatAgentでは、ユーザーが標準のOpenAIまたはAnthropic SDKを通じてループを所有します。Moltは両方のワイヤプロトコルを話すループバックサーバーを起動し、すべてのリクエストはサーバー側でトークン正確な蓄積にデコードされます。長期間動作するエージェントがコンテキストを圧縮してプレフィックスを書き換えると、サーバーは現在のセグメントを自動的に封印し、新しいセグメントを開きます。

3つの正当性不変条件が設計を支えています。トークン同一性(サンプリングされたトークンIDが軌跡を定義し、再トークン化された文字起こしではない)、ポリシーバージョン意味論(訓練可能なトークンは行動ポリシーの対数確率を保持し、非同期使用はシーケンスレベルのゲートの背後でトークンごとに修正される)、前方一貫性(ロールアウトとアクターがモデル意味論について一致する必要がある)です。混合専門家(MoE)ポリシーでは、最後の不変条件が最も重要です。ロールアウトとトレーニングのルーターは独立して専門家を選択するため、小さな数値の差がtop-k選択を反転させる可能性があります。Moltはロールアウトルーティングリプレイを採用し、vLLMがトークンごとの専門家IDを返し、トレーニングの前方計算がそれをリプレイします。

MoltはApache 2.0ライセンスのエージェントRLフレームワークで、RLコードは約8.6K行、verlの約7分の1です。Ray、vLLM、NVIDIA AutoModelをフォークせずに組み合わせ、アップストリームのリリースはコンテナピンとして取り込みます。エージェントは通常のPythonで、標準のOpenAI/Anthropic SDKがそのまま訓練できます。スループットはMegatronベースのスタックと統計的に同等で、MoEの不整合に関する注意点が開示されています。スケールは単なるフラグです。同じループで、--fsdp.ep_size 256を使用して稠密な4Bモデルと700B MoEを実行できます。