ZeRO (Zero Redundancy Optimizer)（ゼロ リダンダンシー オプティマイザー）
Microsoft DeepSpeed が提案したデータ並列学習の冗長メモリを排除する最適化手法。オプティマイザ状態・勾配・パラメータを GPU 間で分割（Sharding）し、各 GPU のメモリ使用量を 1/N に削減する。Stage 1〜3 の段階的な分割レベルがあり、FSDP の理論的基盤でもある。
ZeRO とは
ZeRO（Zero Redundancy Optimizer）は、2019 年に Microsoft Research の Rajbhandari らにより発表された分散学習のメモリ最適化手法です。標準的なデータ並列学習では、全 GPU がモデルの完全なコピー（パラメータ、勾配、オプティマイザ状態）を保持するため、GPU 数が増えてもモデルサイズの上限は 1 GPU のメモリに制約されます。ZeRO はこの冗長性を排除し、各 GPU が全体の一部分のみを保持することで、GPU 数に比例してメモリ制約を緩和します。
3 つの Stage
ZeRO は分割する対象によって 3 つの Stage に分かれます。
| Stage | 分割対象 | メモリ削減（N GPU） | 通信量 | 代表的な用途 |
|---|---|---|---|---|
| Stage 1 | オプティマイザ状態 | 約 4× | 標準 DP と同等 | 中規模モデル |
| Stage 2 | + 勾配 | 約 8× | 標準 DP と同等 | 大規模モデル |
| Stage 3 | + パラメータ | 最大 N× | 約 1.5× | 超大規模モデル |
Stage 1（P_os）: Adam のモメンタムと分散の 2 つの FP32 テンソルを N GPU で分割。7B モデルの場合、8 GPU で各 GPU のオプティマイザメモリが 56GB → 7GB に削減。
Stage 2（P_os+g）: Stage 1 に加え、勾配も分割。各 GPU は自分が担当するパラメータの勾配のみを蓄積。Reduce-Scatter で必要な勾配を集約。
Stage 3（P_os+g+p）: パラメータ自体も分割。順伝播・逆伝播で必要なパラメータをその都度 All-Gather で収集し、計算後に破棄。メモリ効率は最大だが通信量が増加。
通信パターンの詳細
各 Stage での通信オーバーヘッドを理解することが性能チューニングの鍵です。
Stage 1-2: 標準データ並列と同じ通信量。逆伝播中に Reduce-Scatter で勾配を集約し、各 GPU が担当パラメータの更新に必要な勾配を受け取る。その後 All-Gather で更新済みパラメータを全 GPU に配布。
Stage 3: 順伝播の各レイヤーで All-Gather を実行してパラメータを一時的に全 GPU に展開。計算後にメモリから解放。逆伝播でも同様の All-Gather が必要で、合計通信量は Stage 1-2 の約 1.5 倍。ただし通信と計算のオーバーラップ（パイプライニング）で実効オーバーヘッドを軽減。
PyTorch FSDP との関係
PyTorch 2.0 で導入された FSDP（Fully Sharded Data Parallel）は、ZeRO Stage 3 と同等のアルゴリズムを PyTorch ネイティブで実装したものです。
| 比較項目 | DeepSpeed ZeRO | PyTorch FSDP |
|---|---|---|
| エコシステム | DeepSpeed 専用 | PyTorch 標準 |
| 設定方法 | JSON 設定ファイル | Python API |
| オフロード | ZeRO-Infinity（CPU/NVMe） | CPU Offload |
| Mixed Precision | 独自実装 | PyTorch AMP 統合 |
| TP/PP 連携 | Megatron-DeepSpeed | 独自実装が必要 |
DeepSpeed は豊富な最適化オプション（ZeRO-Infinity, 通信最適化など）を提供し、PyTorch FSDP は PyTorch エコシステムとの統合が強みです。
具体的なメモリ計算例
70B パラメータモデル（Llama 2 70B クラス）を 8 × A100 80GB で学習する場合:
ZeRO なし（標準 DP）: 各 GPU にモデル全体を複製
- パラメータ（FP16）: 140GB → 1 GPU に収まらない → 学習不可
ZeRO Stage 3: 全てを 8 分割
- パラメータ: 140GB / 8 = 17.5GB/GPU
- オプティマイザ: 560GB / 8 = 70GB/GPU → 計 87.5GB/GPU → まだ 80GB を超過
ZeRO Stage 3 + Mixed Precision + Activation Checkpointing:
- パラメータ（FP16 分割）: 17.5GB/GPU
- オプティマイザ（FP32 分割）: 35GB/GPU
- 活性化値（チェックポイント）: 約 10GB/GPU
- 合計: 約 62.5GB/GPU → 80GB に収まり学習可能
FAQ
Q1: ZeRO Stage 3 は Stage 1 より常に優れていますか？
A1: メモリ効率は Stage 3 が最善ですが、通信オーバーヘッドが 50% 増加するため、モデルが Stage 1-2 で収まる場合はそちらが高速です。一般的に「収まる最低の Stage」を選ぶのが最適戦略です。
Q2: 個人環境の単一 GPU でも ZeRO は意味がありますか？
A2: ZeRO 自体は複数 GPU 間の冗長性排除が本質なので、単一 GPU では直接的なメリットはありません。ただし ZeRO-Offload は単一 GPU でも CPU メモリにオプティマイザ状態をオフロードでき、単一 RTX 4090 で 13B モデルの学習を可能にします。
Q3: ZeRO と Tensor Parallelism/Pipeline Parallelism は併用できますか？
A3: はい。Megatron-DeepSpeed は ZeRO + TP + PP を組み合わせた「3D 並列」を実現し、千億〜兆パラメータ規模のモデル学習に使用されます。GPT-3 175B や Megatron-Turing NLG 530B はこのアプローチで学習されました。