Transformers、HSDPとExpert Parallelismの統合を提案　MoEの勾配同期を調整

Hugging Face Transformersのメンバーは11月3日、Hybrid Sharded Data Parallel（HSDP）とExpert Parallelismの併用を試みるPR #50293を提出した。この草案では、dp_replicate_size > 1かつep_size > 1の場合に実行を拒否するチェックを削除し、専門家パラメーターのシャーディングと勾配処理も調整している。確認時点で提案はDraftのままで、マージ先は開発ブランチのhsdpとなっており、別の提案にも依存している。正式版でサポートされたと見なすことはできない。出典：開発提案

この取り組みの背景には、複数の並列化戦略でデバイスグループをどう共有するかという課題がある。PyTorchの公式チュートリアルでは、HSDPは2次元構成として説明されている。ホスト内ではFSDPによるシャーディングを行い、ホスト間ではデータ並列レプリケーションを行う。DeviceMeshは各次元の通信グループを管理し、異なる並列化戦略向けにサブメッシュを切り出すこともできる。通常のモデル層と専門家層を同時に配置する学習コードでは、デバイス構成を勾配同期の範囲と一致させる必要がある。出典：PyTorchチュートリアル

コミットの説明によると、データ並列レプリカ数が1を超える場合、トークンの振り分けを担う専門家は(dp_replicate, efsdp)の2次元メッシュでシャーディングされ、通常の密なモジュールは(dp_replicate, fsdp)を使用する。専門家の勾配にはレプリカ間でall-reduceを行い、dp_sizeで除算する。作者は、レプリカの追加に伴って計算に参加するバッチ数も増えるため、それに応じたスケーリングが必要だと説明し、この方法はtorchtitanの実装と一致すると述べている。出典：コミットの説明

設計からは、MoE学習で専門家の分業とデータレプリカによるスケールアウトを組み合わせられる可能性がうかがえる。ただし、実際のスループット、メモリー要件、数値的一貫性は測定が必要だ。公開提案にはまだ性能ベンチマークがなく、ページ上ではレビューも行われていない。エンジニアリングチームは、ベースブランチがマージされるかを追跡し、レプリカ数、Expert Parallelismの度合い、勾配累積の設定を変えて結果を検証する必要がある。特に、スケーリングによって想定どおりの実効バッチサイズの意味が保たれるかを確認したい。出典：提案の状態