Editing Openai/692aba85-5b10-8006-817a-d5a83b020f8e (section)

===== 2.1 prime-rl：大規模非同期RLフレームワーク =====

prime-rl は 「トレーナ」「オーケストレータ」「推論サービス」 の3役で構成されます（図2, p5）。INTELLECT_3_Technical_Report

====== 2.1.1 アーキテクチャ ======
* Trainer - 役割：ロールアウト＋報酬から勾配を計算し、ポリシーを更新 - 実装：PyTorch FSDP2 ベース、torchtitan 由来のコードを使い - データ並列 - テンソル並列 - コンテキスト並列 - MoE向け grouped GEMM - HF 互換モデルをそのまま扱える
* Inference サービス - vLLM をバックエンドにした OpenAI互換 API サーバ - 標準 API に加え /update_weights /reload_weights を追加して、 トレーナからの重み更新を受け取れるようにしている。
* Orchestrator（CPUプロセス） - 推論サーバからロールアウトを受け取り、バッチ化して Trainer へ送る - Trainer から新しい重みを受け取って Inference へ配布 - ロールアウト生成には verifiers 環境を使い、Environments Hub 上の任意の環境をそのまま使える

この3者が疎結合で動作するため、トレーニング用GPUと推論用GPUを完全に分離しつつ、両者を非同期に回せるのがポイントです。

====== 2.1.2 非同期オフポリシー学習 ======
* 通常の同期 on-policy RL だと、 「ロールアウト生成 → 勾配計算 → パラメータ更新」が逐次になるため、 推論側がパラメータ更新待ちで止まる。
* prime-rl では オフポリシーを許容して非同期化： - 推論側は少し古いパラメータ θₙ を使い続けてロールアウトを生成 - その間に Trainer 側は θₙ から θₙ₊₁ へ更新を進める
* 図3（p6）のように、理想化すれば「常に1ステップ遅れ」くらいの off-policy で走るイメージ。

このとき生じる trainer-inference のミスマッチは後述の 重要度サンプリング＋マスキング で制御します。INTELLECT_3_Technical_Report

====== 2.1.3 Continuous Batching & In-Flight Weight Updates ======

長いチェイン・オブ・ソートやエージェントロールアウトでは、ロールアウト長の分散が大きいため、同期バッチだと一番遅いサンプル待ちでGPUが遊びがちです。

そこで、AReal や PipelineRL で広まった手法を取り入れて：
* Continuous Batching - オーケストレータが常に多数のロールアウトリクエストを「プール」しておき、 - ひとつ終わるたび即座に新リクエストを投入 → 推論GPUの利用率を常に高く維持
* In-Flight Weight Updates - Trainer から新しい重みが届くと、推論サーバは一瞬生成を中断して重みをアップデート - その後のトークンは新しいポリシーで生成 - 1本の軌跡の中に複数世代のポリシーが混在し得る - ポリシーの古さが一定ステップ数（max_off_policy_steps）を超えた軌跡は破棄

図4（p7）では、タイムライン上で「ロールアウト開始・終了」と「ポリシー更新」が重なっているイメージが描かれています。INTELLECT_3_Technical_Report

====== 2.1.4 Multi-Client Orchestrator ======

vLLM 標準の multi-node data parallel では、ノード数を増やしてもスループットが頭打ちになったため、
* 各ノードを独立した推論サーバとみなし
* オーケストレータが「ノードごとに1クライアント」を持ち、
* ロールアウトをラウンドロビンで振り分ける

という カスタム data-parallel 方式を採用。これによりノード数にほぼ線形でスケールしたと述べています。INTELLECT_3_Technical_Report

====== 2.1.5 オンライン難易度フィルタ ======

カリキュラム学習として、
* 問題ごとに過去の解答成功率から easy / normal / hard のプールを作成
* 各ステップでプールごとのサンプリング割合を制御
* さらに「常に成功 or 常に失敗」する問題のロールアウトは破棄

といったオンラインフィルタで、常にちょうど良い難易度帯を出し続けるように調整しています。INTELLECT_3_Technical_Report

====== 2.1.6 長いシーケンス長への対応 ======

RL中に自然とシーケンス長が伸びる（長く考えるようになる）ため、48k〜72kトークン級での訓練を効率的に回す必要があります。INTELLECT_3_Technical_Report

アプローチは2つ：
# Context Parallelism（Ring Attention系） - Q/K/V をGPU間で回しながら注意を計算 - 256kトークンまで伸ばせたが、 - データ並列度が半減 - 精度劣化も見られた → 本番設定には不採用
# Activation Offloading - 層出力だけ保存し、中間活性はバックプロップ時に再計算（フルチェックポイント） - さらに出力アクティベーションをCPUへオフロード（torchtuneベース） - これで 48k → 72k まで伸ばしつつ、MFU低下は 0.1% 程度に抑えた

====== 2.1.7 Distributed Muon ======

GLM 系はプレトレで Muon optimizer を使っているため、ポストトレーニングでも Muon を継続使用するのがよいとされる一方、Muon は 行列単位の更新 を行うので、FSDP でシャーディングされた勾配にはそのまま適用できません。INTELLECT_3_Technical_Report

試した方法：
* ラウンドロビン gather/scatter - 各 rank が一部の行列を集めて Newton-Schulz を計算し、更新済み勾配を scatter - しかしノード数が増えると gather が多すぎて InfiniBand が詰まる
* All-to-all ベースの手法（最終採用） - all-to-all collective で勾配を再配置し、Mu on を分散適用 - 実装には Dion の OSS 実装を利用

====== 2.1.8 Mixture-of-Experts の効率化 ======

MoE 層には torchtitan の実装を使い、grouped GEMM カーネルで専門家をまとめて実行。
* 実験の結果、 - シーケンス長と hidden dim が大きい今回の設定では、 既に grouped GEMM が十分に飽和しており、 - Expert Parallelism（EP）を入れてもスループットがむしろ下がることが分かった（scatter/gather オーバーヘッドのみ増える）。INTELLECT_3_Technical_Report
* 図5（p9）のベンチマークでは、 シーケンス長 32k〜65k では 128 experts まではTFLOPSが飽和領域にとどまるため、EPは不要という判断。

推論側では HF の MoE 実装を使う必要があるため、トレーナ側の state dict をHF形式にオンザフライ変換して vLLM に渡す工夫もしています。