rl-triton、7種類の強化学習リターン計算をGPU scanに統合し、関数呼び出し全体で最大5.70倍高速化

rl-tritonは、強化学習のトレーニングで見過ごされがちな信用割り当てフェーズを対象に、Tritonで実装した融合GPU kernel群を公開した。このライブラリは、GAE、V-Trace、Retrace(λ)、TD(λ) returns、discounted returns、eligibility traces、episodic prefix sumsをカバーする。作者によると、これら7種類のアルゴリズムはすべて、一次線形再帰 `A_t = α_t + β_t A_{t+1}` として表現でき、同一の結合演算を用いてassociative scanへ変換することで、逐次実行されていたO(T)の依存チェーンをO(log T)の並列ステージへ置き換えられる。

アルゴリズム専用kernelは、rewards、values、done flags、importance ratiosからレジスタ内でαとβを直接生成し、中間scanをオンチップで完了させる。これにより、doublingの各ステージで結果をHBMへ書き戻す処理を回避する。また実装では、episode termination、時間制限によるtruncation、rolloutウィンドウ境界を区別し、bootstrap値が誤ってゼロに設定されるのを防いでいる。こうした詳細は、特にGAEやoff-policy手法の数値的正確性にとって重要だ。

作者はH100 80GBとRTX 2000 Adaを使用し、4,096個の並列環境、各128ステップという条件でv0.1.3をテストした。すでにO(log T) scanを使用している `torch.compile` ベースラインと比較すると、7種類の演算における関数呼び出し全体の高速化は1.6～5.70倍だった。最高値はRTX 2000 Ada上のdiscounted returnsで記録され、H100上のGAEは2.36倍だった。未コンパイルの逐次ループと比較するよりも、参考価値の高い結果といえる。

ただし、kernelのマイクロベンチマークはトレーニング全体のスループットと同義ではない。ポリシーの隠れ層が1024×1024の場合、GAEがPPO更新に占める割合は約2.2%にすぎず、end-to-endの改善は約1.02倍にとどまる。1.11～1.16倍に達するのは、128×128の小規模ポリシーの場合のみだ。さらに現行版はFP32にしか対応しておらず、系列長4,096のRetraceはレジスタ圧力によりコンパイル済みベースラインより低速になる。また、2種類のforward scanには131,072ステップを超えた場合のfallbackもない。エンジニアリングチームは、まず自社のトレーニングにおいて信用割り当てが占める割合を測定したうえで、統合するかどうかを判断すべきだ。