IsoExec: SkyRLにおけるトレーナーと推論のミスマッチを解消する

vLLMは、強化学習（RL）ワークロードにおけるトレーニングエンジンと推論エンジンの間の数値的なミスマッチを解消するために設計された、クロスフレームワークの統合実行抽象化であるIsoExecを導入しました。厳格な実行契約（execution contract）を強制し、並列性不変（parallelism-invariant）なカーネルを利用することで、IsoExecはロールアウトエンジンとトレーナーがビット単位で一致する同一のポリシーを評価することを保証し、25%のパフォーマンスオーバーヘッドで、ロールアウト対トレーニングの平均エンドツーエンドlogprob差を$10^{-7}$以下に低減します。

問題点：トレーナーと推論のミスマッチ

オンポリシーRLにおいて、システムは同じポリシーを2回実行する必要があります。1回はトークンをサンプリングするためのロールアウトエンジン、もう1回は対数確率を再計算するためのトレーナーです。実際には、これら2つのステージは異なるエンジン（例：ロールアウトにはvLLM、トレーニングにはMegatron）を使用することが多く、それらは異なるカーネル、バッチ形状、および並列レイアウトを採用しています。

浮動小数点演算は結合法則が成り立たないため、これらのシステム上の違いは異なるリダクション順序をもたらし、モデルパラメータが同一であってもトークンの確率を変化させてしまいます。このミスマッチは、REINFORCEやGRPOのようなRLアルゴリズムを不安定化させ、報酬の崩壊（reward collapse）を引き起こす可能性があります。例えば、Fireworksが報告したGLM-5.2の実行では、トレーニングと推論のKLが約0.013であり、その結果、45%のトークンがクリッピングによって破棄され、ステップ20付近で報酬の崩壊が発生しました。

IsoExec 統合実行契約

IsoExecは、両方のランタイムが従うべきすべてのビットに関連する実行の選択を宣言する**実行契約（execution contract）**を使用して、これらの不一致を解決します。このフレームワークに依存しない契約は、丸め誤差に敏感な詳細がエンジン間で同一であることを保証します。

契約の構造と強制

Cases and Regions: 契約はフォワード演算子を「リージョン（regions）」（1つのカーネルによって実装される算術スパン）に分割します。次に、各リージョンがどのように処理されるかを指定するために「ケース（cases）」（例：engine_prefill vs. trainer_fwd）を定義します。
Composition: すべてのリージョンとケースのペアに対して、契約は特定のインプリメンテーションを選択し、蓄積dtypeやリダクション分解パラメータ（例：split-K分割数）などの定数に固定します。
Claims: これらはランタイムによって強制される条件であり、例えば、リダクションツリーがビット単位で不変であることが証明されている特定の並列サイズをリストしたトポロジー・クレーム（topology claims）などがあります。
Identities: SHA-256ダイジェスト（semantic, numerical_policy, および deployment）は、トレーナーとロールアウトエンジンの両方が、検証済みの同一の数値的ポリシーを実行していることを検証するために使用されます。

各ランタイム用の**契約アダプター（contract adapter）**は、これらの仕様をフレームワークの拡張ポイントにバインドし、ランタイムを監視してコンプライアンスを確保します。

統合モデルと並列性不変カーネル

IsoExecは、テンソル並列、エキスパート並列、およびシーケンス並列を含む、さまざまな分散戦略においてビット単位で一致を維持する統合モデル定義を実装しています。

並列性不変性の実現

異なるGPUレイアウトにおいて数値を維持するために、IsoExecは固定バイナリ・リダクションツリー方式を採用しています。

Tensor Parallelism (TP): pikの実装を使用し、IsoExecはK次元を連続するリーフに分割します。各リーフは、FP32蓄積を用いた決定論的なTensor Core MMAを使用し、契約はランクからリーフへのマッピングと算術スケジュールを固定します。
Expert Parallelism (EP): エキスパートの出力は、ランク順ではなく、固定されたルーティング順序で結合されます。
Sequence Parallelism (SP): IsoExecは、非SPシステムのときと同じリダクションツリーを再利用し、各ランクが独自の出力スライスを保持することで、SPが有効であってもトレーナーのロジットが同一であることを保証します。

Chunkwise-Parallel Recurrent (CPR) Gated DeltaNet

Linear-attentionアーキテクチャであるGated DeltaNet (GDN) は、通常、トレーニング/プリフィル用としてチャンク並列形式を用い、デコード用としてリカレント形式を用います。これが数値的なミスマッチを生み出します。これをリカレント形式ですべてを解決しようとする以前の試みは、大幅な速度低下（一部のワークロードで最大5倍）を招きました。

IsoExecは、**Chunkwise-Parallel Recurrent (CPR)**を導入しました。これは、リカレント性をメインの機能として維持しつつ、チャンクを横断して並列に評価します。デコード時には、毎$C$トークンごとに隠れ状態を同期（resynchronize）します（ここで$C$はチャンクサイズです）。これにより、高いスループットを維持しながらビット単位の正確性を実現します。

Stage	Native Mixed	Chunkwise Everywhere	Recurrent Everywhere	CPR
Bitwise Exact	No	Yes	Yes	Yes
Trainer Fwd+Bwd	5.177 ms	5.177 ms (1.00x)	22.863 ms (4.42x)	7.386 ms (1.43x)
Rollout Prefill	0.844 ms	0.844 ms (1.00x)	3.639 ms (4.11x)	1.412 ms (1.67x)
Rollout Decode	0.0612 ms	2.2374 ms (36.6x)	0.0612 ms (1.00x)	0.0846 ms (1.38x)

実験結果

IsoExecは、同期RLを用いたDAPO-Math-17kにおけるQwen3.5-35B-A3Bのトレーニングを、単一の8xH100ノードでテストしました。

数値的精度

50ステップの間に、更新前ロールアウト対トレーニングの平均絶対logprob差は、$10^{-3}$（ネイティブSkyRL）から$10^{-7}$以下（IsoExec）に低減されました。ステップごとの平均最大差は5.073から$10^{-4}$に低下しました。

パフォーマンス・オーバーヘッド

ミスマッチを解消することで、ネイティブSkyRLスタックと比較してパフォーマンス・ペナルティが発生しました。

Metric	Native	IsoExec	Overhead
Generation	591.3 s	776.6 s	31.3%
Policy Training	498.6 s	591.3 s	18.6%
Full RL Step	1224.6 s	1534.0 s	25.3%

数値的整合性は達成されましたが、チームは短い50ステップのテスト期間内では、有意な報酬の改善は見られませんでした。