PyTorch 2.14のインプレース転置で数値誤りの報告、CUDAコンパイル時の融合が関係

9月22日、PyTorchコミュニティで、2.14.0のCUDAコンパイル経路において、一部のインプレース行列転置で数値に誤りが生じると報告された。再現環境はRTX 4070とCUDA 13.0。直接実行した場合は正常だが、`torch.compile`を適用するとテンソルサイズによって結果に差異が生じる。現時点では未解決の公開報告となっている。[問題の記録](https://github.com/pytorch/pytorch/issues/198280)

再現コードは、最後の2つの次元を転置して1を掛けた後、`copy_`で元のテンソルに書き戻す。報告者が8万個の5×5行列でテストしたところ、5回の実行それぞれで、数千から1万超の要素に差異が生じた。その分析では、依存関係チェックとループの並べ替えの相互作用が指摘されている。コンパイラはまず読み書きの順序制約を1つ取り除き、次にループを並べ替え、最後に転置の読み取りとインプレースの書き込みを同じカーネルに融合するという。[再現コードとコンパイルの分析](https://github.com/pytorch/pytorch/issues/198280)

これは、コンパイル後もプログラムの意味が保たれるかどうかに関わる。公式ドキュメントによると、`torch.compile`はデフォルトでInductorを使用し、入力条件に応じて異なるコンパイル結果を生成することがある。したがって、エンジニアリング上の推論として、小さいサイズのケースでテストに通っても、実際のバッチで正しい結果が得られる保証にはならない。また、プログラムが例外を投げないことを確認するだけでは、出力の数値を検証できない。[コンパイルAPI](https://docs.pytorch.org/docs/2.14/generated/torch.compile.html)

公式ドキュメントでは、`clone`は入力のコピーを作成する関数と定義され、デフォルトでは所定の条件を満たすテンソルのメモリ形式を保持する。エンジニアリング上の検証では、コンパイル後の数値がコピーの意味論に沿っているかを確認する必要があり、ソースコードでコピー関数を呼び出しているかどうかだけでは判断できない。[コピーAPI](https://docs.pytorch.org/docs/2.14/generated/torch.clone.html)

報告では、`clone()`や`contiguous()`を使うように書き換えても同じ問題が発生しており、そのまま修正策として扱うことはできない。ループの並べ替えを無効にすると報告者のケースは正常に戻ったが、この内部設定を広く有効なデプロイ時の対策とみなすことは、現段階ではできない。[報告の詳細](https://github.com/pytorch/pytorch/issues/198280)

調査時には、公式ガイドに従ってeager、aot_eager、Inductorの各バックエンドを使い、問題の発生箇所を絞り込める。そのうえで、複数の形状と繰り返し実行の結果を、コンパイルしていないバージョンと比較する。同じバッファを繰り返し更新するワークフローでは、最終出力の形状だけを確認するのではなく、中間テンソルも比較し、バージョン、ドライバー、最小再現コードを保存することを本記事では推奨する。これらは原因を特定するための手段であり、他のGPUやデータ型、モデルへの影響については、引き続き実測が必要だ。[公式トラブルシューティングガイド](https://docs.pytorch.org/docs/2.14/user_guide/torch_compiler/torch.compiler_troubleshooting.html)

同日作成された2.14.1の追跡ページでは、サイレントな数値誤りが、受け入れ可能な重要修正の対象に含まれている。ただし、個々のコミットには申請、テスト、承認が必要とされる。確認時点では、これを根拠に本件が修正済み、あるいは次のバージョンへの修正取り込みが確定したとは言えない。今後は修正コミット、回帰テスト、正式なリリース記録を追跡する必要がある。[バージョンの追跡](https://github.com/pytorch/pytorch/issues/198239)