Prime Intellect、1兆パラメータのMoEモデルをエージェント型RLワークロードで訓練するためのprime-rl 0.6.0をリリース

Prime Intellectは、1兆パラメータのMixture-of-Experts（MoE）モデル向けの強化学習フレームワークprime-rlバージョン0.6.0をリリースしました。このフレームワークは、長期間にわたるソフトウェアエンジニアリングタスクなどの重いエージェント型ワークロードに焦点を当てています。研究チームは、SWEタスクにおいて最大131kのシーケンス長でGLM-5を訓練し、ステップ時間は5分未満、バッチサイズは256ロールアウト、使用したノードはわずか28台のH200でした。

非同期RLの役割 エージェントタスクには長時間実行される外れ値が存在します。一部のコーディングロールアウトは数時間に及ぶことがあり、その都度ポリシー更新を待つとGPUがアイドル状態になります。非同期RLはこれを回避します。トレーナーと推論システムは分離され、独立して実行およびスケーリングされます。ポリシー更新は唯一の同期ポイントであり、prime-rlはオプティマイザのステップが完了するとすぐに新しい重みをプッシュします。既にディスパッチされたロールアウトはアクティブなプレフィックスキャッシュを保持するため、単一のロールアウトが複数のポリシーバージョンのトークンを混合することがあります。

推論最適化 推論は通常、RLシステムにおけるスループットのボトルネックです。prime-rlはスループットを最適化しつつ、レイテンシを制限します。

• FP8推論：DeepEPおよびDeepGEMMカーネルを使用してプリフィルとデコードの精度を低下させます。
• Wide Expert Parallelism：エキスパートを32以上のGPUに分散し、各GPUが個別のエキスパートを保持してエンドポイントとして機能します。同期はレイヤーごとに、ディスパッチとコンバイン操作を通じて行われます。
• プリフィル/デコード分離：プリフィルとデコードのトークン比率が4:1になる場合、共有ワーカーはエンドツーエンドのレイテンシを増加させます。P/D分離により、長いツール出力がデコードワーカーを圧迫するのを防ぎます。
• KVキャッシュ管理：高同時実行性には大きなKVキャッシュスペースが必要です。prime-rlはCPUおよびディスクへの階層的オフローディングをサポートします。vLLMネイティブオフローディングはワーカーごとにプールを作成し、Mooncake Storeは全ノードでRAMとディスクを集中的にプールします。
• ルーターリプレイ（R3）：トレーナーと推論のミスマッチは訓練を静かに破壊します。ルーターリプレイは推論ルーティングの決定をキャプチャし、トレーナーに直接リプレイします。これによりKLミスマッチが約1桁削減されます。

訓練最適化 トレーナーはtorchtitanをベースにしており、3次元並列（FSDP、CP、EP）を採用しています。GLM-5のケーススタディではこれらすべてが使用されました。

• FSDP2：パラメータ、勾配、オプティマイザ状態をシャーディングし、レイヤーごとに重みをオンデマンドで収集します。
• Expert Parallelism：レイヤー内のエキスパートをシャーディングし、アクティブなレイヤーメモリを削減します。all2allディスパッチ/コンバインを使用。
• Context Parallelism：シーケンス次元をシャーディングし、長いコンテキストのアクティベーションメモリを処理します。デフォルトはUlysses、またはRing Attention。
• FP8訓練：DeepGEMMブロックスケーリングFP8を使用し、トレーナーと推論の精度を一致させることでKLミスマッチを減らし、訓練を安定化させます。

ユースケース

• 長期間のSWEエージェント：実際のリポジトリの問題でモデルを訓練。ロールアウトは数百ターンとツール呼び出しに及ぶことがあり、P/D分離によりデコードレイテンシを予測可能に保ちます。
• 少数ノードでの1兆パラメータ規模のポストトレーニング：GLM-5は28台のH200ノードで動作し、Wide EPとKVオフローディングにより同時実行性とスループットを向上させます。
• 大規模で安定したエージェントRL：ルーターリプレイとFP8訓練の両方がトレーナーと推論のKLミスマッチを低減し、より安定した訓練を実現します。

詳細な技術情報は原文を参照してください。