AI-Papers
- 約8.6千行のPyTorch中心コードで、6万行超のverlより大幅に改変しやすく、本番並みの性能を両立した強化学習フレームワーク
- OpenAIやAnthropicのSDKをそのまま学習に使えるトークン基準の設計で、再トークン化によるずれを根本から排除
- Ray、vLLM、NeMoを単一の非同期ループで束ね、700BのMoEまでスケールしつつslimeと同等のスループットを達成
研究の背景
大規模言語モデル（Large Language Model、LLM）をエージェントとして訓練する強化学習（Reinforcement Learning、RL）の研究では、アルゴリズムの改変が日常的に発生します。新しい報酬の推定方法、新しいパイプラインの段階、新しいロールアウト（モデルに実際に文章を生成させて学習データを集める工程）の方式などを、研究者は次々と試していく必要があります。
ところが、広く使われている既存フレームワークでは、こうした小さな変更でも trainer（学習部分）、分散処理のバックエンド、ロールアウトをつなぐ結合コードといった複数の層を貫通して手を入れなければなりません。その負担は、繰り返しのたびに研究者へのしかかります。
実際、代表的なフレームワークである verl はコード規模が約6万2千行に達します。これほどの規模になると、研究者が全体を頭の中で把握するのは難しく、アルゴリズムを少し変えるだけでも大がかりな作業になってしまいます。NVIDIAらの研究チームは、この「改変コスト」の問題に正面から取り組みました。
Moltの全体像
提案手法であるMoltは、PyTorchを中心に書かれた強化学習の訓練フレームワークです。設計の中心にあるのは「研究者が頭の中に収められ、AIコーディングアシスタントが全体を読んで理解できるほど、小さく整理されたコード」という考え方です。
Moltの強化学習部分はわずか約8.6千行のPythonコードで実装されています。参考までに、slimeが約2万5千行、verlが約6万2千行であることと比べると、その簡潔さがよくわかります。コードが小さいと、アルゴリズムの流れを最初から最後まで追いかけて変更しやすくなります。
図1に示すように、Moltは3つの部品と1つのループで構成されます。3つの部品とは、素のPythonで書かれたユーザーエージェント、リクエストルーターの後ろに置かれた複数のvLLMロールアウトエンジン、そしてNeMo AutoModel上で動く単一のFSDP2ポリシーアクター（実際に学習される本体）です。これらを1つのRayの非同期キューが統合しています。
エージェントを特別な仕組みではなく普通のプログラムとして扱える点が、Moltの設計思想を端的に表しています。エージェントの振る舞いを自由に組み替えられるこの発想は、自己改善ループで探索能力を高めるリサーチエージェントの研究とも問題意識が重なります。
単一の非同期ループ
Moltは「streaming pool（ストリーミングプール）」と呼ばれる仕組みで非同期処理を実現します。これは複数のプロンプト群を同時に保持し続ける、いわば作業中の待ち行列です。十分な数のプロンプト群が生成を終えた時点でバッチが組まれるため、学習が動いている間もエンジンが手持ち無沙汰になりません。
特徴的なのが「partial rollout（部分ロールアウト）」です。モデルの重みを更新するとき、生成の途中だったリクエストを捨てる必要がありません。Moltはエンジンを一時停止し、アクターのパラメータを配り直したうえで、保持していたリクエストをそのまま再開します。
再開したリクエストには、更新前と更新後で異なるバージョンのポリシーが混在します。この問題に対して、Moltはすべての行動トークンがサンプルされた時点のlog確率（生成のしやすさを表す値）を保持し、トークンごとの補正を損失計算に適用することで、数値的な整合性を保っています。
重みの配り直し（weight refit）は、リクエストルーターを経由しません。アクターから各ロールアウトエンジンへ、NCCL（GPU間の高速通信ライブラリ）を使ってパラメータを直接ブロードキャストします。各エンジンは自分が担当する部分だけを読み込み、ルーターは推論のトラフィックだけを運ぶため、通信のボトルネックが生じにくくなります。
3つの正しさ保証
Moltが重視するのは、学習の「正しさ」です。生成していないトークンで学習しないという原則のもと、以下の3つの不変条件を保証します。
- トークンの一致: 再トークン化した文章ではなく、実際にサンプルされたトークンIDが軌跡を定義する
- ポリシーバージョンの整合: 学習対象のトークンは生成時のlog確率を保持し、非同期利用のずれを補正する
- 順伝播の一致: マルチモーダル展開やMoEのルーティングを含め、ロールアウトと学習でモデルの意味づけを揃える
特に3つ目は、専門家混合モデル（Mixture of Experts、MoE）で重要になります。MoEは入力ごとに使う専門家（部分ネットワーク）を選び分けますが、生成時と学習時でその選び方がずれると計算が食い違ってしまいます。Moltは「ルーティングのリプレイ」で生成時の専門家の選択を記録し、両者を一致させています。
SDKをそのまま使える設計
Moltのもう1つの強みは、OpenAIやAnthropicの標準SDKをそのまま学習に使える点です。仕組みとしては、エンジンの前にloopback（折り返し）のチャットサーバーを立て、そこを通るリクエストをサーバー側で1つのトークン単位の記録へと変換します。これをTITO（token-in/token-out）キャプチャと呼びます。
この方式では、トークンの空間から一度も外に出ないため、再トークン化によるずれが発生しません。エージェント側のコードには特別な引数もセッション管理の追加も不要で、外部のツールや自動化の仕組みをそのままポリシーの駆動に使えます。研究の再現性とツールの流用しやすさを両立する、実用的な工夫といえます。
実験結果
主要な評価には、Qwen3-30B-A3Bを2ノード（学習用8GPU、ロールアウト用8GPU）で用い、コンテキスト長16,384トークン、応答上限8,192トークンという設定が使われました。さらにマルチモーダルのMoEや、700BのMoEをexpert parallelism 256でスケールさせる超大規模の検証も行われています。
性能面では、Megatronベースの最先端スタックであるslimeと真っ向から比較されました。結果は次の通りで、両者は統計的に同等の水準にあります。
指標 | Molt | slime |
|---|---|---|
最適化1ステップの時間 | 119.4 ± 2.3秒 | 109.5 ± 10.3秒 |
スループット | 461 トークン/GPU/秒 | 502 トークン/GPU/秒 |
RL部分のコード行数 | 約8.6千行 | 約2万5千行 |
コード量が3分の1程度でありながら、本番並みのスループットを維持している点がMoltの成果です。エンジン最適化の効果も報告されており、prefix caching（先頭部分の再利用）ではキャッシュ命中時の再計算が0.05秒に短縮され、投機的デコーディングでは生成時間が329秒から64秒へと約5倍高速化しました。
一方で、最適化器の状態をCPUへ退避する設定では、GPUメモリを64.7GBから46.4GBに削減できる代わりに約18%の速度低下が生じています。メモリと速度のどちらを取るかは、用途に応じた選択になります。
まとめと今後の展望
Moltは、エージェント強化学習の研究で繰り返されるアルゴリズム改変のコストを、コードの簡潔さで下げる試みです。約8.6千行のPyTorch中心コードで、6万行超のverlより格段に手を入れやすく、それでいてslimeと同等の本番並み性能を保っています。
トークン基準の設計による3つの正しさ保証、単一の非同期ループによる部分ロールアウト、700B MoEまでのスケール実証など、実用と研究の双方に配慮した設計が特徴です。Apache-2.0ライセンスで全公開されており、レシピやコンテナもNeMo Labsのリポジトリから利用できるため、再現性の高さも評価できます。
ただし、slimeとの比較では最適化1ステップの時間やスループットでわずかに劣っており、絶対的な速度で上回るわけではありません。あくまで「改変しやすさと本番性能の両立」に価値がある手法です。コードの読みやすさが研究の反復速度をどこまで押し上げるか、今後の活用事例が注目されます。
