Transformers、特定構成での専門家並列勾配の二重除算を修正

Hugging Face Transformers は11月5日、PR #50292 をメインブランチにマージし、特定の専門家並列構成で混合専門家モデル（MoE）の勾配スケーリングが誤る問題を修正しました。この問題は token dispatch と FSDP2 の相互作用に関係しています。メンテナーの説明によると、efsdp_size == 1 の場合、専門家の勾配が二重に除算され、最終的に想定より fsdp_size 倍小さくなります。修正提案

この条件は並列トポロジーに関係します。公式ドキュメントでは efsdp_size = fsdp_size * tp_size / ep_size と定義されており、専門家全体の複製数を表します。たとえば、テンソル並列を無効にし、fsdp_size=8、ep_size=8 とした場合、専門家は8つのプロセスに分散され、efsdp_size は1になります。Token dispatch では各プロセスが独自のバッチを保持し、2回の all-to-all によって token を対応する専門家に送り、結果を受け取ります。専門家並列のドキュメント

この不具合は勾配リダクションの境界条件で発生します。通常の経路では、まず FSDP2 が勾配を合計し、その後データ並列の規模に応じて除算します。しかし、専門家側の FSDP シャードが1つしかない場合、実行できる reduce-scatter がありません。FSDP2 は勾配を複製する時点ですでに除算しているにもかかわらず、強制的な合計設定によってさらに除算していました。修正では、reduce-scatter が必要な場合に限って合計を強制するようにし、専門家の勾配が一度だけスケーリングされるようにしています。メンテナーによる根本原因の説明

学習チームにとって、この問題は専門家パラメーターと他のモジュールとの更新比率に影響する可能性があるため、並列構成を比較する際に調査する価値があります。これはスケーリングの仕組みに基づく工学的な推論です。公開提案には収束曲線や下流タスクの品質低下に関する完全なデータが示されていません。そのため、勾配が8分の1になったことをモデル品質が8分の1になったと解釈することはできず、すべての MoE 学習に影響すると推論することもできません。

エンジニアはまず、実際に使っている専門家の分配方式と、上記3つの並列パラメーターを確認できます。そのうえで、データを固定して修正前後の専門家の勾配ノルムと1ステップの更新量を比較してください。公式ドキュメントには、メインブランチの内容はソースコードのバージョンに対応しており、token dispatch を使用する場合のバッチおよび token 数のカウントに関する Trainer の対応は今後補完される予定であることも記されています。今回のマージで解決したのは特定の勾配の問題です。導入前に、インストール済みバージョンに修正が含まれているか確認し、学習パイプライン全体を検証してください。公式の使用上の制限