AI生成CUDA kernel、benchmark通過後にtrainingを壊したbf16の罠

r/MachineLearningで、AI生成CUDA kernelの難しい失敗例が注目を集めた。NVIDIAのSOL-ExecBenchで上位に入ったkernelを実際のproductionに近いworkloadへ入れたところ、いくつかは意外な形で壊れたという。問題になった一つは、fused embedding-gradientとRMSNorm backward passを処理するkernelで、benchmark verifierは通過したが、小さなtransformerのtrainingではlossが発散した。

原因は単純な誤答ではなかった。embedding-gradient側の累積がfp32ではなくbf16で行われていた。tokenを均等にサンプリングすると寄与が分散し、bf16でも問題が見えにくい。実際のテキストでは一部のtoken IDに多数のgradientが集まり、小さな値が大きくなったaccumulatorに丸め込まれて消える。その結果、高頻度のembedding rowが少しずつdriftする。AdamWではper-parameter normalizationがこのbiasを吸収し、loss上では問題が隠れた。

コメント欄の論点は「verifierを通った」という言葉の弱さだった。bf16はよく使われるため見落としやすいという指摘や、optimizerとdataset sensitivityをkernel testに含めるべきだという意見が出た。厄介なのは、この種のbugが研究アイデアの失敗、datasetの問題、architectureの問題に見えてしまう点だ。

AIが生成するperformance codeは十分に速くなっている。次の課題は、実際のtrainingやservingで静かに間違えないかを検証することだ。速度を競うbenchmarkには、より広い分布とoptimizer条件を含む検証が必要になる。

Reddit discussion · Related research post