FSDP（エフエスディーピー）
Fully Sharded Data Parallel の略。PyTorch が提供する分散学習フレームワークで、モデルパラメータ・勾配・オプティマイザ状態を全 GPU にシャーディングし、メモリ効率を最大化しながらデータ並列学習を行う。
FSDP とは
FSDP（Fully Sharded Data Parallel）は、PyTorch 1.11 で導入され PyTorch 2.x で大幅に改良された分散学習 API である。Microsoft DeepSpeed の ZeRO-3 と同等の機能を PyTorch ネイティブで提供し、モデルパラメータ・勾配・オプティマイザ状態をすべて GPU 間でシャーディング（分割配置）する。Meta が Llama 2/3 のファインチューニングに使用したことで広く認知された。
なぜ FSDP が必要か
従来の DDP（DistributedDataParallel）では、各 GPU にモデル全体のコピーを保持する必要があった。70B パラメータモデルの場合、FP16 パラメータだけで 140GB、オプティマイザ状態（Adam の m, v）を含めると約 560GB が各 GPU に必要となり、物理的に不可能である。
FSDP は N 台の GPU にパラメータを 1/N ずつ分割して保持し、計算時のみ AllGather で一時的に全パラメータを収集する。これにより GPU あたりのメモリ消費を劇的に削減する。
| 方式 | GPU あたりメモリ（70B FP16 + Adam） | 8 GPU 構成 |
|---|---|---|
| DDP | 560 GB | 不可能 |
| FSDP | 70 GB | 可能（H100 80GB） |
FSDP の動作原理
Forward Pass
- 各 GPU は自分のシャード（1/N のパラメータ）のみ保持
- 各レイヤーの計算前に AllGather で全パラメータを収集
- Forward 計算を実行
- 計算後、収集したパラメータを解放（メモリ節約）
Backward Pass
- AllGather でパラメータを再収集
- Backward 計算で勾配を算出
- ReduceScatter で勾配をシャーディング（各 GPU が 1/N の勾配を受け取る）
- 収集したパラメータを解放
パラメータ更新
各 GPU が保持するシャードのパラメータとオプティマイザ状態のみを更新。通信なしで独立に実行できる。
FSDP2（PyTorch 2.x）
PyTorch 2.x で torch.distributed.fsdp が刷新され、FSDP2 と呼ばれる新 API が導入された。主な改善点:
- DTensor ベース: テンソル並列（TP）との統合が容易に
- per-parameter sharding: レイヤーごとにシャーディング戦略を柔軟に設定
- 混合精度の改善: BF16 Forward + FP32 パラメータの切り替えがシームレスに
- コンパイラ対応:
torch.compileとの互換性向上
シャーディング戦略
FSDP では用途に応じて複数のシャーディング戦略を選択できる。
| 戦略 | 内容 | メモリ効率 | 通信量 |
|---|---|---|---|
| FULL_SHARD | パラメータ+勾配+オプティマイザをシャード | 最高 | 最大 |
| SHARD_GRAD_OP | 勾配+オプティマイザをシャード（パラメータは保持） | 中 | 中 |
| NO_SHARD | シャードなし（DDP 相当） | 低 | 最小 |
FULL_SHARD が最もメモリ効率が良いが、AllGather/ReduceScatter の通信オーバーヘッドが最大。GPU 間の帯域が十分な場合は FULL_SHARD、帯域が制約される場合は SHARD_GRAD_OP が適する。
FSDP の実装パターン
PyTorch での FSDP 実装は比較的簡潔である:
model = FSDP(
model,
sharding_strategy=ShardingStrategy.FULL_SHARD,
mixed_precision=MixedPrecision(
param_dtype=torch.bfloat16,
reduce_dtype=torch.float32,
),
auto_wrap_policy=size_based_auto_wrap_policy,
)
auto_wrap_policy がレイヤーの粒度でシャーディング単位を自動決定する。Transformer モデルでは transformer_auto_wrap_policy を使い、各 TransformerBlock 単位でラップするのが標準パターンである。
FSDP vs DeepSpeed ZeRO
| 観点 | FSDP | DeepSpeed ZeRO |
|---|---|---|
| エコシステム | PyTorch ネイティブ | 独自ライブラリ |
| ZeRO ステージ | ZeRO-3 相当 | ZeRO-1/2/3/Infinity |
| CPU オフロード | 対応（限定的） | ZeRO-Infinity で NVMe まで |
| TP 統合 | DTensor で対応 | Megatron-DeepSpeed |
| チェックポイント | 分散チェックポイント | 独自形式 |
Meta や Hugging Face のエコシステムでは FSDP が標準、Microsoft 系や大規模プレトレーニングでは DeepSpeed が優勢という住み分けが見られる。
FAQ
Q1: FSDP を使うべきモデルサイズの目安は？
A1: 単一 GPU のメモリに収まらないモデル（おおむね 7B パラメータ以上のフルファインチューニング）で威力を発揮する。7B 未満や LoRA ファインチューニングなら DDP で十分なことが多い。
Q2: FSDP と LoRA は併用できるか？
A2: 可能だが、LoRA のパラメータ数が少ないため FSDP のメリットは薄い。LoRA で単一 GPU に収まるなら DDP が効率的。FSDP + QLoRA の組み合わせは、70B モデルのファインチューニングで消費者向け GPU（RTX 4090 等）を複数使うシナリオで有用。
Q3: FSDP のチェックポイント保存は通常と異なるか？
A3: FSDP ではパラメータがシャーディングされているため、保存時に全パラメータを 1 箇所に集約する必要がある。PyTorch 2.x の DistributedCheckpoint API を使えば、シャーディングされたまま保存・復元でき、チェックポイントのI/O 時間を大幅に短縮できる。