PyTorchコミュニティがコンパイルキャッシュの正確性に関する問題を報告、計算グラフを保持した再度の逆伝播で勾配が誤る可能性

PyTorchコミュニティは10月25日、issue #201970を提出し、特定のコンパイルキャッシュ状態で torch.compile を使うと、保持した計算グラフに対する2回目の逆伝播で、例外が発生せずに誤った勾配が生成される可能性があると報告しました。確認時点でこのissueは未解決で、修正PRは記載されていません。投稿者は現象を再現できたと述べていますが、原因の診断が正しいかどうかはまだ確信がないとしています。issue報告

最小の再現例では、まず二乗関数をコンパイルして通常の逆伝播を実行し、次にDynamoをリセットして再コンパイルした後、retain_graph=True を指定して勾配を2回計算します。報告によると、連続して4になるはずの結果が4、8になります。一方、直接実行モードやキャッシュが空の状態では、2回とも4になります。記事中のAIによる分析では、上書き可能な中間バッファーがキャッシュキーで区別されていないため、後続の呼び出しで適切でない逆伝播用コードが再利用されている可能性が疑われています。この仕組みは、まだ検証されていない仮説です。再現コードと分析

公式ドキュメントによると、FXGraphCacheとAOTAutogradCacheはコンパイル結果をディスクに保存できるため、別のプロセスでも同じグラフと設定に遭遇すれば、再コンパイルを省けます。両者は指定されたキャッシュのルートディレクトリも共有します。そのため、実務上はキャッシュの状態と実行順序も確認する必要があります。初回コンパイルだけをテストしても十分とは限りません。これは、ドキュメントと報告内容に基づく検証上の提案です。キャッシュ設定のドキュメント

もう一つの確認点は勾配の累積です。公式APIの説明では、backward は葉ノードの勾配を累積します。上記のケースでは勾配を計算するたびに .grad をクリアしているため、報告された数値の差を通常の累積動作だけで説明することはできません。また、ドキュメントでは元の計算グラフを保持する retain_graph と、導関数の計算グラフを作成する create_graph が区別されています。テストではそれぞれを個別に記録する必要があります。逆伝播のドキュメント

計算グラフを保持する機能とコンパイル済みトレーニングを利用するチームは、まず直接実行、キャッシュが空の状態、キャッシュを再利用する状態の3通りで結果を比較できます。今後は、メンテナーによる独立した再現確認、影響を受けるバージョン、回帰テストに注目すべきです。現時点の証拠だけでは、一般的なトレーニング処理に広く影響すると結論づけることはできません。