[論文レビュー] PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel

Yanli Zhao, Andrew Gu+16 more|arXiv (Cornell University)|Apr 21, 2023

Advanced Neural Network Applications被引用数 12

ひとこと要約

この論文は PyTorch Fully Sharded Data Parallel (FSDP) を提示します。パラメータシャーディングによる大規模モデル訓練のためのネイティブで業界グレードのソリューションであり、遅延初期化、構成可能なシャーディング、そして広範な最適化を備え、小さなモデルでは DDP と同等の性能を発揮し、はるかに大きなモデルではほぼ線形の TFLOPS スケーラビリティを実現します（512×80GB A100 GPU まででテスト済み）。

ABSTRACT

It is widely acknowledged that large models have the potential to deliver superior performance across a broad range of domains. Despite the remarkable progress made in the field of machine learning systems research, which has enabled the development and exploration of large models, such abilities remain confined to a small group of advanced users and industry leaders, resulting in an implicit technical barrier for the wider community to access and leverage these technologies. In this paper, we introduce PyTorch Fully Sharded Data Parallel (FSDP) as an industry-grade solution for large model training. FSDP has been closely co-designed with several key PyTorch core components including Tensor implementation, dispatcher system, and CUDA memory caching allocator, to provide non-intrusive user experiences and high training efficiency. Additionally, FSDP natively incorporates a range of techniques and settings to optimize resource utilization across a variety of hardware configurations. The experimental results demonstrate that FSDP is capable of achieving comparable performance to Distributed Data Parallel while providing support for significantly larger models with near-linear scalability in terms of TFLOPS.

研究の動機と目的

• 単一GPUメモリを超えた非常に大規模モデルの訓練に対する業界グレードでフレームワーク内蔵の解決策の必要性を動機づける。
• PyTorch FSDP を紹介し、非侵襲で効率的な訓練を提供するための PyTorch コアコンポーネントとの共設計を示す。
• ハードウェアの異種性と大容量メモリフットプリントを扱うためのシャーディング戦略、初期化方法、メモリ・通信最適化を開発・評価する。
• 大規模GPUクラスター上で言語モデルと推奨モデルの両方に対して FSDP の有効性とスケーラビリティを示す。

提案手法

• モデルをFSDPユニットに分解し、1つのユニットのみを完全に材料化するようパラメータをシャードする。
• 遅延初期化を実装し、偽デバイス上でモデル構築を許し、実際のGPU上で材料化時に初期化をリプレイする。
• シャーディング係数Fを用いて、メモリフットプリントと通信のトレードオフを図る完全シャーディングとハイブリッドシャーディングの複数戦略を提供する。
• FlatParameter を用いてパラメータを統合し、AllGather/ReduceScatter の追加コピーを避けつつ効率的に行う。
• ビューと勾配フックを用いた PyTorch autograd への密接な統合により、正確で自動的な勾配伝播とリダクションを実現する。
• 通信最適化（オーバーラップ、フォワード/バックワードのプレフェッチ、勾配蓄積）とメモリ管理（CUDA キャッシュアロケータ、レートリミッター）を適用する。

実験結果

リサーチクエスチョン

• RQ1 異種ハードウェア全体でスループットを最大化しつつメモリを最小化するためのネイティブ PyTorch ソリューション設計はどうあるべきか。
• RQ2 小規模モデルでは DDP に匹敵する性能を達成し、はるかに大きなモデルではほぼ線形の TFLOPS スケーラビリティを提供できるか。
• RQ3 単一GPUのメモリを超える訓練モデルをサポートするために必要な初期化、シャーディング、メモリ管理技術は何か。
• RQ4 完全シャーディング setup で計算と重ね合わせを最大化するために通信をどのように組織すべきか。
• RQ5 実データセンターのトポロジーでの完全シャーディングにおける完全シャーディングとハイブリッドシャーディングのトレードオフは何か。

主な発見

• FSDP は小規模モデルで DDP の性能に匹敵しつつ、はるかに大きなモデルを可能にし、ほぼ線形の TFLOPS スケーラビリティを実現できる。
• FlatParameter ベースのシャーディングアプローチにより、集約/分散操作中のコピーオーバーヘッドを最小限に抑えつつ、パラメータレベルのメモリ効率を実現。
• 遅延初期化とユニット単位の材料化により、事前の全メモリ割り当てなしで非常に大きなモデルを構築できる。
• ハイブリッドシャーディングはホスト間のトラフィックを削減し、データセンターのローカリティを活用できるため、中規模モデルに対してメモリとスループットの可調整トレードオフを提供する。
• FSDP は PyTorch 2.0 のベータ版であり、産業界と研究環境での戦闘テストが行われ、再利用可能な成果物が提供されている。
• 実験は大規模 GPU クラスターでのスケーラビリティを示しており、最大 512 x 80GB A100 GPU にまで達している。