PyTorchコミュニティが係数の乗算漏れを報告、2.14のCUDAモメンタム更新で値が10倍になる可能性

PyTorchのCUDAコンパイル経路で、オプティマイザーの数値を気づかれないまま変える可能性のある回帰事例が見つかった。ユーザーのrwightmanは10月5日、timmのAdafactorモメンタム更新を処理する際、torch.compileで実行したPyTorch 2.14の結果が大きくなると報告した。添付された最小再現例では、コンパイル後の出力は直接実行した場合の10倍だった。同じケースで、2.13のCUDA経路と2.14のCPU経路では正しい比率が得られた。問題報告

この挙動を引き起こす処理は、まず更新テンソルを二乗平均平方根に基づいてクリッピングし、次にスケーリング係数を乗じ、最後にm.mul_(0.9).add_(u, alpha=0.1)を実行するというものだ。公式APIの定義では、alphaは加算されるテンソルに乗じられるため、この処理は0.9m + 0.1uを計算するはずだ。しかし、報告に添付されたTritonコードでは前半のスケーリングだけが保持され、最後の0.1が抜けていた。テストでモメンタムをゼロに初期化すると、10倍の差が生じる。この倍率を、モメンタムがゼロでない場合や学習全体の結果にそのまま当てはめることはできない。公式の加算ドキュメント、インプレース加算のドキュメント

10月6日に提出された修正案は、Inductorのパターンマッチングに対処の焦点を当てている。ノードに、パターンで宣言されていない既定値以外のキーワード引数がある場合は置換を適用しないようにし、融合演算で引数の意味が失われるのを防ぐ。ただし、提案者は説明文がAIによって生成され、本人の確認をまだ受けていないと明記している。その後のボットレビューでもテストと説明に問題があると指摘されており、この案を検証済みの修正とみなすことはできない。修正提案

コンパイル済みオプティマイザーを扱うエンジニアリングチームにとって、この事例はスループットだけでなく、更新テンソルとモメンタム状態の値も確認する必要があることを示している。報告では、lerp_への書き換えで正しい結果が得られたが、これはこの再現例で観察された結果に限られる。今後は修正案のレビュー、回帰テスト、正式リリースへの収録状況を追い、自社のオプティマイザーでもコンパイル前後の数値を比較すべきだ。現時点で公開されている証拠だけでは、影響を受けるGPU、データ型、長時間学習の範囲は判断できない。