要注意:AI生成のCUDAカーネルが学習と推論を密かに破壊する現象が発生
先月、NVIDIAはDeepSeekやQwen、Gemma、Kimiなどから抽出した235個の本番環境向けCUDAカーネルを含む新しいベンチマーク「SOL-ExecBench」を公開しました。私たちは、このベンチマークで高スコアを記録したAI生成コードをいくつか試しに本番ワークロードへ導入してみたのですが、驚くべきことにその多くが意図せず機能停止を引き起こしました。その一つが、Transformerの学習ステップの最後で実行される「fused embedding-gradient + RMSNorm backward pass」です。ベンチマークの検証を余裕でクリアした最速のカーネルを採用したところ、学習の損失(Loss)が発散し、二度と回復しなくなりました。原因究明のためデータセットやオプティマイザ(SGDからAdamWへ変更)を入れ替えると症状が消えるという、研究において最も厄介なパターンに陥りました。これは「アイデア自体がダメだったのか?」という疑念を生み、研究者がデータやアーキテクチャの修正に無駄な時間を費やしてしまう種類のバグです。調査の結果、実際のバグは「カーネルの勾配計算の一部がfp32ではなくbf16で蓄積されていたこと」でした。Embedding backwardは小さな勾配を積み重ねますが、実データでは特定のトークンIDにアクセスが集中するため、bf16では精度不足となり値がゼロに丸められて漂流してしまうのです。AdamWはそのバイアスを吸収するため一見問題がないように見えますが、他のアルゴリズムでは致命的なバグとなります。他の壊れた実装例についても、私たちのブログ記事で詳細に解説しています。