NVIDIA の NeMo チームは、珍しい設計目標を持つエージェント強化学習用の PyTorch ネイティブ フレームワークである Molt をリリースしました。それは、研究者が頭の中に保持できるほど十分に小さなコードベースであり、AI コーディング アシスタントがその全体を読み取って推論できるというものです。このリリースは、エージェント RL がツールの使用、コードの記述、環境のナビゲートをモデルに教えるための主要なレシピとなる中で導入されました。これは、フロンティア エージェントのトレーニング方法を再構築するいくつかのインフラストラクチャ プロジェクトの 1 つです。これらの変化については、AI ニュース速報 で取り上げています。

実行するだけでなく、読み取れるように構築

MarkTechPost が報告しているように、Molt は各フレームワークの RL エントリ ポイントからインポート グラフをトレースすることによってカウントされ、約 8.6K 行の RL コードを測定します。同じ方法で、verl の場合は約 62K 行、スライムの場合は 25K 行、OpenRLHF の場合は 7.2K 行が集計されます。その意図的なコンパクトさがプロジェクトの中心的なピッチです。エージェント的 RL の研究は、アルゴリズムの継続的な変更 (新しい推定器、新しいパイプライン ステージ、新しいロールアウト スキーム) であり、主流のフレームワークでは、すべての変更がトレーナー、分散バックエンド、ロールアウト接着剤の層を通過します。 Molt はその摩擦を取り除くことを目指しています。

このフレームワークは Apache 2.0 ライセンスであり、起動コード、Slurm スクリプト、および構築済みコンテナが付属しています。しかし、NVIDIA は、添付の研究論文で Molt を実稼働トレーニング サービスではなく研究インフラストラクチャとして位置付けており、ハードウェアが本当の門番であることを明確にしています。出荷されたレシピは、8 つの H100 GPU の 2 ノード、トレーニング用に 8 つ、ロールアウト用に 8 つを分割することを前提としています。これにより、フロンティアおよびフロンティアに隣接する研究室、ポストトレーニングを行っている資金豊富なスタートアップ企業、エンタープライズ AI 研究グループ、およびマルチノード H100 または H200 アクセスを持つ学術グループが利用できるようになります。

フォークではなく構成済み

Molt のアーキテクチャは、フォークすることなく 3 つの既存のツールを構成するため、アップストリームの改善は面倒なリベースではなく、コンテナのピンとして提供されます。配置と非同期キューには Ray、ロールアウトには vLLM、トレーニングには NVIDIA AutoModel with FSDP2 が使用されます。ランタイムは、エージェント プール、リクエスト ルーターの背後にある一連の vLLM エンジン、および単一のトレーニング可能なポリシー アクターで構成されます。ストリーミング プールはプロンプト グループを飛行状態に保つため、アクターのトレーニング中にエンジンが空になることはありません。

部分ロールアウトと呼ばれる機能は、効率化の中心となります。ポリシーが更新されると、Molt はエンジンを一時停止し、アクターの更新されたシャードを NCCL 経由で各エンジンに直接ブロードキャストし、保持されたリクエストを破棄せずに再開します。これにより、モデルが変更されるたびに進行中のロールアウトを破棄するという無駄なパターンを回避できます。

1 つのモジュール、2 つのエージェント フォーム

Molt で実行される RL は、AgentRunner をエクスポートする単一の Python モジュールに名前を付けます。報酬関数を含むその他のすべては通常のコードです。フレームワークは 2 つの形式をサポートします。 Env を使用すると、フレームワークは Gymnasium に合わせた `step()` 内に LLM ループを所有します。これは、おなじみの強化学習パターンに適合します。 ChatAgent を使用すると、ユーザーはストック OpenAI または Anthropic SDK を通じてループを所有し、既存のエージェントを簡単にラップできるようになります。

両方の橋渡しをするために、Molt は両方のワイヤ プロトコルを話す ループバック サーバー を起動し、すべてのリクエストはサーバー側でデコードされて 1 つのトークンとまったく同じ蓄積になります。長期エージェントがコンテキストを圧縮してプレフィックスを書き換えるとき (多くのターンにわたって実行されるエージェントの一般的な操作)、サーバーは現在のセグメントをシールし、新しいセグメントを自動的に開きます。これは、エージェント RL の実行が実際に正しいか、それとも単に正しく見えるかどうかを決定する一種の配管の詳細です。

3 つの正確性の不変条件

Molt の設計は、非同期エージェント トレーニングの微妙な失敗に対処する 3 つの正確性不変条件を中心に構成されています。 トークン ID とは、再トークン化されたトランスクリプトではなく、サンプリングされたトークン ID が軌道を定義することを意味します。テキストをトークンにつなぎ戻すと、黙ってデータが変更される可能性があるため、これが重要です。 ポリシー バージョン セマンティクス により、トレーニング可能なトークンが動作ポリシーのログ確率を維持し、シーケンス レベルのゲートの背後で非同期使用がトークンごとに修正されることが保証されます。 前方一貫性 では、ロールアウト エンジンとトレーニング アクターがモデルのセマンティクスに同意する必要があります。

この最後の不変条件は、ますます一般的になっている 専門家混合 (MoE) 政策 にとって最も重要です。ロールアウト ルーターとトレーニング ルーターは専門家を個別に選択します。わずかな数値の違いにより、上位 k 個の選択肢が反転し、サンプリングされたものとトレーニングされたものの間に不一致が生じる可能性があります。 Molt は ロールアウト ルーティング リプレイ でこれに対処します。vLLM はトークンごとのエキスパート ID を返し、トレーニング フォワード パスはそれらの正確なルーティング決定を再導出するのではなくリプレイします。

用途

同梱されたレシピとユースケースは、NVIDIA が Molt の採用を期待している場所、つまりマルチターンのツール使用エージェント、コード実行エージェント、ビジョン言語環境 (フレームワークには同梱の geo3k レシピが含まれています)、LLM-as-judge 報酬ループ、および小規模な学生モデルへのポリシー蒸留を示しています。これらはまさに、業界全体でエージェント RL の急増を引き起こしたワークロードであり、実際のトレーニングが課す部分的なロールアウト、非同期サンプリングの条件下で適切に処理するのが難しいことで知られています。

より広範なシグナルは、NVIDIA がエージェントをトレーニングする GPU だけでなく、研究者がエージェントを構築するために使用するソフトウェア スタックにも投資しているということです。コードベースを小さくて読みやすいものに保ち、AI コーディング アシスタントがコードベースを変更できるように明示的に設計することで、Molt は、将来のエージェント RL ツールがそれを使用するモデルと共同開発されるという賭けを反映しています。

AI の一歩先を行く

フロンティア エージェントのトレーニング方法を再構築するフレームワークの詳細については、最新の AI 開発 のハブをフォローしてください。

AIニュースをもっと読む→