GitHub Repo
PyTorchコミュニティ、BF16コンパイル時の勾配のずれを報告 畳み込み後のlog_softmaxが学習に影響する可能性
10月17日のPyTorchコミュニティの報告によると、BF16混合精度で特定の畳み込みとlog_softmaxを含む計算グラフをInductorでコンパイルすると、誤った勾配が生じる可能性がある。報告には再現コードと暫定的な回避策が添えられているが、メンテナーは影響範囲をまだ確認しておらず、修正版も発表していない。
PyTorchコミュニティは10月17日、学習の正確性に関する問題を報告した。CUDA BF16 autocastを有効にした状態で、畳み込みにlog_softmaxを続ける計算グラフに後続の演算を加えると、torch.compileによるコンパイル後に勾配がずれる可能性がある。報告環境はPyTorch 2.10.0、CUDA 12.8。確認時点でissueは未解決で、関連する修正もない。問題報告
再現コードでは、本来同じ勾配になるはずの2つの関数を比較している。一方にはdetach()を適用した項が追加されている。未コンパイル時は両者の勾配が一致するが、コンパイル後は、logitsへの共通オフセットを1000にすると勾配のコサイン類似度が0.6943まで低下し、ノルム比は2.295に達した。これにより、学習曲線から異常を推測するだけでなく、開発者が検証できる具体的な事例が示された。再現データ
報告者による生成カーネルの分析では、順方向の演算はFP32のlogitsを使って最大値とlog-sum-expを計算する一方、logitsをBF16バッファーに格納する。逆方向の再計算では、丸め済みのlogitsと、丸め前のデータに基づく統計量が混在し、順方向と逆方向の計算に不整合が生じる。これはあくまで報告者による根本原因の見立てであり、メンテナーの検証はまだ行われていない。カーネル分析
公式の数値精度ドキュメントによると、浮動小数点演算の順序やバックエンドの違いによって結果が変わることがあり、数学的に等価な計算でもビット単位の一致は保証されない。そのため、エンジニアリングチームはずれの大きさと収束への影響を調べ、未コンパイルの経路も比較対象として残しておくとよい。数値精度ドキュメント
報告者によると、aot_eagerに切り替えるか、torch._inductor.config.emulate_precision_casts = Trueを設定すると、報告されたケースでは問題を回避できる。同様の演算チェーンを使うチームは、まずコンパイル前後で勾配の方向とノルムを比較できる。この内部設定による性能上のコストは、各自で測定する必要がある。今後は、メンテナーによる再現の確認、影響を受けるバージョン、修正が回帰テストに含まれるかを追跡したい。暫定対応と今後の確認事項