要注意：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はそのバイアスを吸収するため一見問題がないように見えますが、他のアルゴリズムでは致命的なバグとなります。他の壊れた実装例についても、私たちのブログ記事で詳細に解説しています。

AdamWでは動くのにSGDでは動かないからといって、bf16の代わりにfp32を使うのがバグだとは思えないな。

結局、本当のバグはカーネルの埋め込み勾配（embedding-gradient）の半分がfp32ではなくbf16で累積されていたことだった。

うわ、それは多くの人が絶対に見つけられないような類のものだね。bf16はよく使われるから、それを見て見過ごしちゃう人もいそうだ。

じゃあ、解決策はAdamWを使うことだったんだね。

うわ。一体どうやってそんなバグが起きたんだ？fp32が必要な箇所をbf16で置き換えちゃったってこと？

投稿のソースコードはどうやって見るの？