PyTorch、DTensorの数値修正を提案　NaN処理による分散リダクション結果の変化を防止

PyTorchコミュニティは11月7日、分散テンソルDTensorの演算配置ルールを修正するPR #203933を提出した。nan_to_numは、リダクション待ちの最大値・最小値の配置を維持できる演算として扱われていたが、NaNやカスタムの無限大置換値がある場合、この最適化によって結果が変わる可能性がある。確認時点で提案はオープンのままで、まだマージされていない。修正提案

DTensorは配置を使ってデータの分散方法を表す。Partialは各プロセスがリダクション前のローカル値を保持している状態で、Replicateは各プロセスが完全な複製を保持している状態を示す。公式ドキュメントによると、前者から後者への変換ではall_reduceが実行される。そのため、配置を維持するルールは通信と演算の順序も決める。DTensorドキュメント

2つのプロセスがそれぞれNaNと5を持つ例で問題を説明できる。報告された通常のテンソルの意味に従えば、先に最大値を取ると結果はNaNとなり、その後nan_to_numで置換すると0になる。一方、各プロセスで先にNaNを0に置き換えてから最大値を取ると、結果は5になる。カスタムのposinf、neginf置換値でも、演算順序の等価性が損なわれる可能性がある。これは演算の意味の違いであり、浮動小数点の許容誤差ではない。問題と再現コード

この修正では、3つのnan_to_numオーバーロードに関する配置維持ルールを削除する。これにより、システムは先にReplicate()へ変換してリダクションを完了し、その後に特殊値を置換する。投稿者によると、CPUのLocalTensorModeとtorch 2.14.1の環境で4つのケースを再現し、修正後はいずれも期待どおりの結果となった。また、関連する要素単位のテストにも合格したという。ただし、これは投稿者の報告であり、ページ上では関連CIの実行はまだ承認待ちとなっている。検証記録

エンジニアリングチームは、シャーディングされたテンソルに対して最大値または最小値のリダクションを行った直後に、NaNやInfを処理している箇所を確認するとよい。特殊値を含む入力を使い、分散テンソルと通常のテンソルの結果を比較できる。ドキュメントと修正の仕組みから推測すると、リダクションを早めに完了させることで通信のタイミングも変わる可能性があり、性能は別途測定が必要だ。今後はメンテナーのレビュー、実際のマルチプロセステスト、正式リリースへの収録状況を追う必要がある。現時点の証拠から、一般的な学習フロー全般が影響を受けるとは結論づけられない。