テクニカルガイド
完全にシャード化されたデータ並列
Fully Sharded Data Parallel (FSDP) は、モデルのパラメーター、勾配、オプティマイザーの状態を多くの GPU に分割して、各デバイスがスライスのみを保持する分散トレーニング手法です。
このガイドを聴くAIによるナレーション · ~2 min
概要
これにより、モデル全体を 1 つの GPU のメモリに収めることができないハードウェア上での巨大なモデルのトレーニングが可能になります。
ディープダイブ
従来のデータ並列処理では、モデルの完全なコピーがすべての GPU に保持されるため、メモリが無駄になり、モデル サイズに上限が生じます。 Meta の PyTorch によって普及し、Microsoft の ZeRO からインスピレーションを得た FSDP は、パラメーター、勾配、オプティマイザーの状態という 3 つのものをデバイス間でシャーディングします。フォワード パス中、各 GPU は、オールギャザーを介して計算しているレイヤーの全重みを一時的に収集し、計算を実行した後、収集したコピーをすぐに解放します。バックワード パスも同様に機能し、その後に、勾配スライスを所有する GPU に分配するリデュース スキャッターが続きます。各デバイスはモデルの一部のみを永続的に保存するため、メモリ使用量は GPU の数に応じてほぼ直線的に減少し、チームが数百億または数千億のパラメーターを使用してモデルをトレーニングできるようになります。
現実世界の実装
個別に最大の重みを保持できない 8 つの GPU にわたる 700 億パラメータの Llama モデルを微調整します。
AI ラボで、数百のアクセラレータにわたってオプティマイザーの状態 (Adam でメモリを支配する) をシャーディングすることにより、大規模な言語モデルを事前トレーニングします。
研究者は、PyTorch の FSDP ラッパーを使用して、主力の 80 GB GPU を購入せずに大学クラスターでビジョン トランスフォーマーをトレーニングしています。
FSDP と混合精度 bfloat16 を組み合わせると、メモリが約半分になり、マルチモーダル モデルのトレーニング スループットが高速化されます。
完全にシャード化されたデータ並列の未来
FSDP はオープンな大規模モデルのトレーニングのデフォルトになりつつあり、PyTorch の FSDP2 により使いやすさとパラメーターごとのシャーディングが向上しています。兆パラメータ モデルのテンソルおよびパイプライン並列処理とのより緊密な統合、混合精度と fp8 のサポートの向上、シャーディング境界を選択するよりスマートな自動ラッピングが期待されます。 NVLink や InfiniBand などの GPU 間の相互接続が高速化するにつれて、シャーディングの通信コストは縮小し続け、これまで以上に大規模なスケールで実用的になります。
実践する
なぜそれが重要なのか
- 費用と予算. アーキテクチャの決定により、パフォーマンスと運用コストが何年にもわたって推進されます。
- より明確な判決. 技術教育は、チームが最新のスタックだけでなく、適切なスタックを選択するのに役立ちます。
- 品質管理. より良いエンジニアリングの選択により、本番環境での信頼性に関するインシデントが減少します。
気をつけてください
- 1 つのベンチマークを最適化すると、より広範なシステムの弱点が隠れる可能性があります。
- インフラストラクチャとメンテナンスのコストは過小評価されがちです。
- システムが複雑になるにつれて、セキュリティと可観測性のギャップが拡大する可能性があります。
次のステップ
- 実装前にレイテンシ、品質、コストの目標を定義します。
- 現実的な負荷とデータ条件でのベンチマーク。
- エラー、ドリフト、ユーザーへの影響を計測器で監視します。
- スケーリングの前に、ロールバックとインシデント対応のパスを準備します。
探検を続けましょう
よくある質問
完全にシャード化されたデータ並列とは何ですか?
Fully Sharded Data Parallel (FSDP) は、モデルのパラメーター、勾配、オプティマイザーの状態を多くの GPU に分割して、各デバイスがスライスのみを保持する分散トレーニング手法です。これにより、モデル全体を 1 つの GPU のメモリに収めることができないハードウェア上での巨大なモデルのトレーニングが可能になります。
標準のデータ並列処理では実行できない、FSDP の GPU 間でのシャーディングとは何ですか?
FSDP はモデルのパラメーター、勾配、オプティマイザーの状態をデバイス間でシャード化しますが、標準のデータ並列処理ではすべての GPU で完全なモデルを複製します。
FSDP がレイヤを計算する直前にレイヤの全重みを再構築するために使用する集合演算はどれですか?
レイヤーが実行される前に、FSDP はオールギャザーを実行してすべてのシャードから完全なパラメーターを一時的に組み立て、その後それらを解放します。
FSDP がレイヤーの計算直後に収集された全重みを解放するのはなぜですか?
シャードのみを永続的に保持し、全ウェイトを一時的に収集することで、メモリ使用量が低く抑えられ、モデルの一部にほぼ比例します。
FSDP は主に、初期のメモリ最適化アプローチからインスピレーションを得たのはどれですか?
FSDP のパラメーター、勾配、オプティマイザー状態のシャーディングは、DeepSpeed ライブラリの Microsoft の ZeRO で導入されたアイデアに厳密に従っています。
FSDP は重みの収集からネットワーク遅延の大部分をどのようにして隠しますか?
FSDP は、現在の層がまだ計算している間に次の層のパラメータをプリフェッチし、オールギャザー通信に有用な作業を重ねます。