TransformersのメインブランチにHSDP設定を追加、シャーディングとレプリカを組み合わせた学習に対応

Transformersのメインブランチのコミット履歴に、11月6日付でHSDP（ハイブリッドシャーデッド・データ並列）のサポートが追加され、Trainerの勾配累積に関する調整も含まれました。複数ノードの学習環境を自社で管理するチームにとって、この変更によりモデルの読み込み時に「各レプリカを何枚のGPUに分散するか、レプリカを合計いくつ作るか」を直接指定できるようになります。メインブランチのコミット履歴

追加された設定は DistributedConfig(dp_replicate_size=R, fsdp_size=F) です。パラメータをF台のデバイスにシャーディングし、それをR組に複製します。提案では、2ノードに各2枚のGPUがある構成を例に、両方の値を2に設定し、各ノードに1つずつシャーデッド・レプリカを配置しています。勾配はまずレプリカ内でreduce-scatterを行い、その後レプリカ間でall-reduceを行います。HSDPの提案

これは、既存の分散学習技術をTransformersのモデル読み込みインターフェースに統合するものです。PyTorchのFSDP2では、2次元のDeviceMeshを使って同じ構造を表現できます。一方の次元にレプリカを配置し、もう一方の次元でシャーディングを担います。GPU間の接続やノードのトポロジーに応じて通信グループを構成できることが、エンジニアリング上の利点です。ただし、実際のスループットはモデルサイズ、ネットワーク帯域幅、メモリの余裕に左右されるため、新しい設定の追加だけから高速化の程度を推測することはできません。PyTorchドキュメント

関連するTrainerの変更は、別の統合上の問題に対応しています。DistributedConfigを通じて読み込み時にシャーディングされたモデルはAccelerateでラップされていないため、従来の accelerator.no_sync(model) では各マイクロバッチの同期を止められませんでした。新しい経路では、HSDPの場合、レプリカ内のreduce-scatterを維持しながら、レプリカ間のall-reduceのみを遅らせます。通常のFSDPでは勾配同期を遅らせます。sync_each_batch=True の場合は、引き続きバッチごとに同期します。勾配累積の提案

導入を検討するチームは、まずインストール済みのバージョンに該当コミットが含まれているかを確認し、そのうえで通信時間、ピークメモリ使用量、学習結果を測定してください。公開PRページには、より早い時点の未マージまたはドラフトの状態が表示されており、新しいメインブランチの履歴との間に時間差があります。現時点で確認できるのは、関連コミットがメインブランチに存在することです。これだけでは、正式リリース版で完全なサポートが提供されているとは判断できません。