Command Palette
Search for a command to run...
CUDA-L1:対照強化学習を活用したCUDA最適化の向上
CUDA-L1:対照強化学習を活用したCUDA最適化の向上
Xiaoya Li Xiaofei Sun Albert Wang Jiwei Li Chris Shum
概要
GPU計算リソースへの需要の急増により、自動化されたCUDA最適化戦略の導入が急務となっている。近年の大規模言語モデル(LLM)の進展はコード生成において有望な成果を示しているが、現在の最先端モデルではCUDAの実行速度向上という点で低効率な成果にとどまっている。本論文では、新規の対照型強化学習(contrastive RL)アルゴリズムを採用した自動化強化学習フレームワーク「CUDA-L1」を提案する。
CUDA-L1は、CUDA最適化タスクにおいて顕著な性能向上を達成した。NVIDIA A100上で訓練された本モデルは、KernelBenchに含まれる全250個のCUDAカーネルに対して平均で3.12倍の高速化(中央値:1.42倍)を実現し、ピークでは最大120倍の高速化を達成した。さらに、本モデルはGPUアーキテクチャ間でも高い移植性を示しており、A100専用に最適化されたにもかかわらず、L40では平均3.12倍、RTX 3090では2.50倍、H100では2.39倍、H20では2.37倍の平均高速化を実現した。
CUDA-L1の能力は、初期段階で性能が低かったLLMが、実行速度に基づく報酬信号のみを用いて、人間の専門知識やドメイン知識を一切不要にすることで、有効なCUDA最適化ツールへと進化可能であることを示している。このアプローチは、CUDA演算の自動最適化の実現可能性を広げ、GPUの計算効率を大幅に向上させ、GPUリソースへの需要増加による圧力を緩和する可能性を秘めている。一方で、本研究ではCUDA開発のようなタスクにおける強化学習モデルの訓練に伴う重要な課題も明らかにした。具体的には、報酬関数の仕様に不備がある場合、RLは目的の最適化問題を解くのではなく、報酬関数の穴(ロジックの抜け)を悪用する傾向があることである。こうした失敗モードを特定し、その根本原因を分析することで、報酬の不正利用(reward hacking)を防ぐことができる実用的な訓練プロシージャの構築方法を提案した。