## Memo ## Memo with LLM ### 論文情報 - **論文のタイトル**: PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel - **著者と所属**: Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, Alban Desmaison, Can Balioglu, Pritam Damania, Bernard Nguyen, Geeta Chauhan, Yuchen Hao, Ajit Mathews, Shen Li（全員Meta AI所属） - **カンファレンス/ジャーナル名**: Proceedings of the VLDB Endowment (PVLDB) - **発表年**: 2023年 ### 論文概要 大規模モデルの訓練を民主化するため、PyTorchのコアコンポーネントと密接に連携設計されたPyTorch Fully Sharded Data Parallel (FSDP)を提案している。FSDPは、従来のDistributed Data Parallel (DDP)では扱えない大規模モデルを効率的に訓練可能にし、TFLOPSにおいてほぼ線形のスケーラビリティを実現している。 ### 詳細解説 #### 問題設定 **問題の入力と出力**: 大規模ニューラルネットワークモデルの分散訓練において、単一GPU上のメモリ容量を超えるモデルサイズに対応する必要がある。入力はモデルパラメータ、勾配、オプティマイザ状態であり、出力は効率的な分散訓練の実現である。 **必要なデータ**: モデルパラメータ、訓練データ、複数のGPUデバイス、高速ネットワーク接続が必要である。実験では最大512台の80GB A100 GPUを使用している。 #### 提案手法 **FSDPアルゴリズム**: モデルを小さなユニット（FSDP units）に分解し、各ユニット内のパラメータをフラット化してシャード（shard）する。Forward/Backward計算時のみパラメータを一時的に復元し、計算後即座に破棄することで、メモリ使用量を最小化する。 **FlatParameter**: 各FSDP unit内の全パラメータを1次元テンソルに連結し、パディングを追加してシャーディング因子で割り切れるサイズにする。これにより効率的なAllGatherとReduceScatter操作を実現する。 **シャーディング戦略**: - フルシャーディング（F = W）: 最小メモリ使用量、最大通信オーバーヘッド - ハイブリッドシャーディング（1 < F < W）: メモリとスループットのトレードオフ調整 - フルレプリケーション（F = 1）: DDPと同等 #### 新規性 **先行研究との比較**: ZeRO や DeepSpeed の影響を受けているが、PyTorchのコアコンポーネント（Tensor実装、ディスパッチャシステム、CUDAメモリキャッシュアロケータ）との密接な統合により、よりロバストで一貫したユーザー体験を提供している。 **技術的新規性**: - Deferred initialization: "fake" デバイス上でモデル初期化操作を記録し、実GPU上で再実行 - Communication optimization: 通信と計算のオーバーラップ、backward prefetching、rate limiter - Native mixed precision: シャーディング環境での効率的な混合精度訓練 #### 実験設定 **使用データセット**: - HuggingFace T5-11Bトランスフォーマー - minGPT-175Bトランスフォーマー - DHEN推薦モデル（768B sparse + 550M dense parameters） - WikiHowデータセット（テキスト要約タスク） **評価指標**: - TFLOPS per GPU（スループット指標） - バッチあたりレイテンシ - ピークメモリ使用量（allocated, active, reserved） - GPU使用率 #### 実験結果 **スケーラビリティ**: GPT-175Bモデルで128-512GPU間においてほぼ線形スケーリング（173-186 TFLOPS per GPU）を達成。これはA100の理論ピーク性能（312 TFLOPS BF16）の55-60%に相当する。 **メモリ効率**: T5-11Bモデルにおいて、DDPではOut-of-Memory エラーが発生する規模でもFSDPは正常動作。シャーディング戦略により柔軟なメモリ-スループットトレードオフが可能。 **通信最適化効果**: Backward prefetchingにより約18%の性能向上。Rate limiterはT5モデルで最大5倍の速度向上を実現（ただしモデル依存）。 ## Abstract 大規模モデルが幅広いドメインで優れた性能を発揮する可能性があることは広く認識されている。機械学習システム研究の分野で大きな進歩が遂げられ、大規模モデルの開発と探索が可能になったにもかかわらず、そのような能力は少数の上級ユーザーと業界リーダーに限定されており、より広いコミュニティがこれらの技術にアクセスし活用する上での暗黙の技術的障壁となっている。本論文では、大規模モデル訓練のための産業グレードソリューションとして、PyTorch Fully Sharded Data Parallel (FSDP)を紹介する。FSDPは、Tensor実装、ディスパッチャシステム、CUDAメモリキャッシュアロケータなど、PyTorchのいくつかの主要コアコンポーネントと密接に共同設計されており、非侵入的なユーザー体験と高い訓練効率を提供する。さらに、FSDPは、さまざまなハードウェア構成にわたってリソース使用率を最適化するための一連の技術と設定をネイティブに組み込んでいる。実験結果は、FSDPがDistributed Data Parallelと同等の性能を達成しながら、TFLOPSの観点で準線形スケーラビリティを持つ大幅に大きなモデルのサポートを提供できることを実証している。