GitHub Repo
PyTorchコミュニティが融合SGDの初期化漏れを報告、AMPが初回ステップをスキップするとモメンタム状態が汚染される可能性
10月21日のPyTorchコミュニティからの報告によると、勾配オーバーフローで初回の更新がスキップされた場合、融合SGDに未初期化のモメンタムバッファが残り、次の有効な更新に影響する可能性があります。別の参加者はCPUで再現したと述べていますが、メンテナーによる確認や正式な修正バージョンはまだありません。
PyTorchでは10月21日、学習の正確性に関する報告が追加されました。torch.optim.SGDでmomentum>0とfused=Trueを同時に設定し、勾配に無限大が含まれるためGradScalerが初回の更新をスキップすると、オプティマイザーに未初期化のモメンタム状態が残り、次の有効な更新時に読み取られる可能性があります。報告者は最小再現コードを提示し、2.14.1と10月20日のnightlyで、CPUとCUDAの両方に問題が発生したと述べています。コミュニティの報告
この相互作用には、混合精度学習の保護処理が関係しています。公式AMPの例では、GradScalerはまず勾配のスケールを元に戻し、infまたはNaNを検出すると更新をスキップして、異常な勾配によるパラメーターの書き換えを防ぐと説明されています。そのため、この種の問題を調べる際は、その時点でパラメーターに変化がないことを確認するだけでなく、後続の計算に影響する状態がオプティマイザーに残っていないかも確認する必要があります。AMPドキュメント
別の参加者はCPUで再現したと述べ、原因は融合パスにあると説明しました。処理は先にモメンタムバッファを確保し、カーネルはオーバーフローフラグを検出すると早期リターンしますが、外側の処理はそのバッファを保持します。すると次のイテレーションで、バッファが既存のモメンタムとして扱われます。報告にある決定論的な設定では未初期化メモリにNaNが入るため、その後パラメーターがNaNになります。これは問題を明らかにするテスト条件であり、すべての学習で必ず同じ結果が起きることを意味するものではありません。再現とソースコード分析
公式のSGDドキュメントでは、初回のモメンタムにはその時点の勾配をそのまま使い、初回にはdampeningを適用しないと定めています。融合版は演算をまとめて性能を高めますが、現在はデフォルトで有効ではありません。したがって修正では、未初期化値の読み取りを防ぐとともに、初回の有効な更新における数学的な意味を保つ必要があります。SGDドキュメント
この設定を使うエンジニアリングチームが検証する際は、初回ステップで意図的にオーバーフローを起こし、fused=Trueとfused=Falseでパラメーターとモメンタム状態を比較する方法が考えられます。今後は、メンテナーによる確認、非ゼロのdampeningを含む回帰テスト、修正がどのバージョンに実際に取り込まれるかを追う必要があります。確認時点では、議論はコミュニティによる再現と修正案にとどまり、修正PRへのリンクはありません。イシューの状況