RL-Align/RL-Kernel

解決する課題

RL-Kernelは、強化学習 (RL) のポストトレーニング中に発生する数値的な不整合に対処します。標準的なセットアップでは、ロールアウトエンジン（サンプル生成）とトレーニングエンジン（モデル更新）が異なるカーネルや計算順序を使用することが多く、これが対数確率の不一致を招き、RLトレーニングの安定性と品質を低下させる原因となります。RL-Kernelは、これら2つのステージで同一の結果を生成することを保証し、トレーニングプロセス全体を通じて数値的一貫性を維持します。

仕組み

実行エンジン（ロールアウト用の vLLM やトレーニング用の Megatron-LM など）と基盤となるハードウェアとの間のミドルウェア層として機能します。attention、dense FFN、および GRPO や PPO といった特定のRL目的関数など、決定論的で高性能な演算子セットを提供することで、ロールアウトエンジンとトレーニングエンジンの両方に数学的に同一の計算を実行させます。また、サポートされているアクセラレータ上でパフォーマンスを最適化し、メモリ使用量を削減するためのハードウェア固有のカーネルも含まれています。

対象ユーザー

このプロジェクトは、LLMの大規模RLポストトレーニングに取り組んでおり、高性能で再現性が高く、クロスハードウェアに対応したトレーニングパイプラインを必要とする機械学習エンジニアや研究者向けに設計されています。

主な特徴

• 決定論的一貫性： ロールアウトエンジンとトレーニングエンジン間でのビット単位で正確な対数確率の一致を保証します。
• クロスハードウェアサポート： NVIDIA (CUDA) および AMD (ROCm) GPUをネイティブサポートし、Huawei Ascendハードウェアも部分的にサポートしています。
• パフォーマンス最適化： 融合計算とカスタムカーネルを使用して、ロールアウト時間と同期オーバーヘッドを最小限に抑えます。
• エンジン統合： vime と統合して vLLM および Megatron-LM をオーケストレーションし、RLワークフローのための統一されたパスを提供します。