PyTorch Fully Sharded Data Parallel を使用した大規模モデル学習の加速

Hugging Face は、PyTorch Fully Sharded Data Parallel (FSDP) を Accelerate ライブラリに統合しました。これにより、オプティマイザの状態、勾配、およびパラメータをデータ並列ワーカー間でシャード化（分割）することで、大幅に大規模なモデルの学習が可能になります。この統合により、より大きなバッチサイズの使用が可能になり、CPU オフロードを通じて、本来であれば GPU メモリの制限を超えるようなモデルの学習も可能になるため、大規模モデルの学習が民主化されます。

FSDP vs. Distributed Data Parallel (DDP)

PyTorch FSDP は、すべての GPU 上にモデルの完全なレプリカを保持することに伴う冗長なメモリ消費を排除することで、Distributed Data Parallel (DDP) を改善します。

DDP では、すべてのワーカーがモデルのパラメータ、勾配、およびオプティマイザの状態の完全なコピーを保持します。各ワーカーは異なるデータバッチを処理しますが、モデルを更新する前に、すべてのワーカー間で勾配を平均化するための all-reduce 操作を行う必要があります。

FSDP では、オプティマイザの状態、勾配、およびモデルのパラメータがワーカー間でシャード化されます。フォワードパスおよびバックワードパスの実行中、FSDP は all-gather 操作を使用して、特定のレイヤーまたはラップされたモジュールに必要なパラメータのみを取得し、計算直後にそれらを解放します。その後、ローカルな勾配は reduce-scatter 操作を介して平均化および分散され、各ワーカーが自身のローカルなパラメータのシャードのみを更新できるようにします。このアプローチにより、GPU ごとのメモリ使用量を劇的に削減できます。

GPT-2 でのパフォーマンス ベンチマーク

Hugging Face は、Causal Language Modeling タスクにおいて、2 枚の NVIDIA Titan RTX GPU (各 24GB) を使用して FSDP と DDP のベンチマークを実施しました。

GPT-2 Large (762M Parameters)

FSDP は、DDP と比較して大幅に大きなバッチサイズを可能にします。CPU オフロードなしの場合、FSDP は最大 15 のバッチサイズを可能にします (DDP の 7 と比較して)。CPU オフロードを有効にすると、バッチサイズはさらに 22 まで増加します。混合精度 (FP16) を使用した DDP が生の学習時間においては最速でしたが、FSDP は、より大きなバッチサイズに必要なメモリ効率を提供します。これは、動的なバッチングを行うアプリケーションにおいて特に有益です。

GPT-2 XL (1.5B Parameters)

GPT-2 XL モデルの場合、DDP はバッチサイズ 1 の時点でも CUDA Out of Memory (OOM) エラーが発生しました。対照的に、FSDP は学習の成功を可能にしました：

• FSDP (Zero-Stage 3): 2 枚の GPU 上で、GPU ごとにバッチサイズ 5 をサポート。
• FSDP with CPU Offload: 単一の GPU でバッチサイズ 10、および 2 枚の GPU で GPU ごとにバッチサイズ 14 の学習を可能にしました。

技術的な実装と設定

Accelerate による統合

ユーザーは、accelerate config CLI を使用するか、あるいは FullyShardedDataParallelPlugin を介してより詳細な制御を行うことができます。主な設定オプションは以下の通りです：

• Sharding Strategy: FULL_SHARD と SHARD_GRAD_OP の間の選択。
• Min Num Params: デフォルトの auto-wrap policy によってレイヤーがラップされるために必要な最小パラメータ数。
• Offload Params: パラメータと勾配を CPU にオフロードするかどうかを決定する boolean 値。

Auto Wrap Policy の役割

min_num_params の設定は、メモリ最適化において極めて重要です。default_auto_wrap_policy を使用する場合、FSDP はパラメータ数が指定された閾値を超えるレイヤーをラップします。BERT-Large (330M) でのベンチマークでは、auto-wrap を使用した FSDP は DDP の約半分のメモリを消費することが示されています。min_num_params を低く設定する (例：2k) ことで、、高い閾値 (例：1M) と比較して、メモリ使用量をわずかにさらに削減できます。

重要な注意点と制限事項

FSDP を使用する実務家は、いくつかの技術的な制約に留意する必要があります：

• Optimizer Initialization: FSDP はパラメータをフラット化し、その場でシャード化します。したがって、モデルはオプティマイザが作成される 前に accelerator.prepare(model) を介して準備される必要があります。オプティマイザをモデルのラップ前に作成すると、オプティマイザが壊れるか、メモリ使用量が増加する可能性があります。
• Parameter Groups: FSDP はネストされたモジュールを 1D 配列にフラット化するため、ラップ前に作成されたパラメータグループ (例：バイアスに異なる weight decay を適用する場合) は、単一のグループに統合され、失われます。
• Multiple Models: 複数のモデルを学習する場合、エラーを避けるために、それぞれのオプティマイザを作成する前にモデルを準備することが必須です。
• Mixed Precision: 本稿の公開時点では、PyTorch の修正待ちのため、FSDP では混合精度がサポートされていませんでした。

分散学習アプローチのまとめ

FSDP は、大規模なモデルを扱うために設計された、より広範な分散学習戦略のエコシステムの一部です：

• ZeRO (Zero Redundancy Optimizer): FSDP の基盤であり、オプティマイザの状態 (Stage 1)、勾配 (Stage 2)、およびパラメータ (Stage 3) をシャード化します。
• Tensor Parallelism: 個々の大きなレイヤーのパラメータを GPU 間でシャード化します。
• Pipeline Parallelism: 異なるレイヤーを異なる GPU に分散させ、マイクロバッチをパイプライン化します。
• 3D Parallelism: ZeRO、Tensor、および Pipeline parallelism の組み合わせであり、数百億のパラメータを持つモデルに使用されます。