Transformersコミュニティ、学習時の正規化に関する修正を提案。有効トークン数の切り捨てで勾配が増幅する可能性

Transformersコミュニティは11月1日、学習時の数値の正確性に関する問題を報告した。Trainerが損失の正規化に使う分母を計算する際、有効なターゲットトークン数がモデル並列係数で割り切れないと、切り捨てによって勾配が増幅される可能性がある。報告された環境は開発版の5.23.0.dev0、Tesla T4 2枚、FSDP2とcp_size=2を組み合わせたコンテキスト並列設定だった。問題報告

問題の箇所は_get_num_items_in_batch()にある。有効トークン数を、テンソル並列、コンテキスト並列、シーケンス並列の各係数の積で割る際に、//を使っている。たとえば有効なターゲットが3つある場合、本来は1.5となる分母が1に切り捨てられる。報告者による最小限の確認では、有効なターゲットが1つの場合、分母が0に切り捨てられることも示された。再現コードとカウント結果

このカウントはログ表示だけでなく、損失計算にも使われる。Trainerの公式ドキュメントによると、カスタム損失インターフェースには勾配累積バッチ全体の項目数が渡され、num_items_in_batchが正規化に使われる。そのため、実際の分母と更新幅を確認する必要があり、学習処理が完了するかどうかだけでは十分ではない。Trainerドキュメント

報告者が提出した修正では、整数除算を通常の除算に置き換え、割り切れるケース、割り切れないケース、分母が0になるケースの回帰テストを追加した。報告者による2 GPUでの検証では、有効なターゲットが3つの場合、修正前の損失、勾配ノルム、1ステップのSGD更新ノルムはいずれも参照実装の1.5倍だった。一方、ターゲットが4つの比較では結果が一致した。これらは提出者自身による報告であり、PRはまだオープンで、レビューも確認されていない。修正提案

コンテキスト並列では系列方向に入力を分割し、複数のGPUで長い系列のAttention計算を分担する。Accelerateのドキュメントも、ラベルのシフトと分割方法が損失の正確性に影響すると注意を促している。今回の報告は、正規化の分母も確認対象になることを示しているが、すべての並列化方式や正式リリース版が影響を受けると結論づける根拠にはならない。今後はメンテナーによる確認、ほかの並列化方式での実機テスト、修正が最終的に含まれるバージョンに注目したい。コンテキスト並列のドキュメント