# ZeROメモリ最適化 ## 定義 ZeRO(Zero Redundancy Optimizer)は、データ並列(DP)訓練におけるメモリ冗長性を排除するシステム最適化フレームワークである。[[Samyam Rajbhandari]]・Jeff Rasley・Olatunji Ruwase・[[Yuxiong He]](Microsoft)が提案し、SC 2020 で発表、[[DeepSpeed]] として OSS 化された。([[@2020__SC__ZeRO Memory Optimizations Toward Training Trillion Parameter Models]]) 設計の核心は「**モデル状態は常時全プロセスに保持する必要はない**」という洞察にある。各状態が必要とされるタイミング(オプティマイザ状態はステップ末尾、勾配は後退伝播直後、パラメータは各レイヤーの計算時)を利用し、状態を分割して保持しつつ必要時のみ集団通信で集約する動的スケジュールを採用する。 ZeRO は 2 系統の最適化からなる。 - **ZeRO-DP**: データ並列のモデル状態(オプティマイザ状態・勾配・パラメータ)の冗長排除 - **ZeRO-R**: 残余メモリ(活性化・一時バッファ・断片化メモリ)の最適化 ### ZeRO-DP の 3 段階 | 段階 | 分割対象 | メモリ削減 | 通信量(DP 基準) | |---|---|---|---| | Stage 1(P_os) | オプティマイザ状態 | 約 4 倍 | 変化なし(2Ψ) | | Stage 2(P_os+g) | +勾配 | 約 8 倍 | 変化なし(2Ψ) | | Stage 3(P_os+g+p) | +パラメータ | Nd 倍 | 1.5 倍(3Ψ) | Stage 1〜2 は**通信量を一切増やさずにメモリを削減できる**ため、既存 DP ワークロードへの移行コストが低い。Stage 3 のみ通信量が 1.5 倍増加するが、それでもモデル並列がノード間に拡張される場合の通信コスト増に比べれば大幅に効率的。 ### ZeRO-R の 3 要素 - **Pa(分割活性化チェックポイント)**: モデル並列が複製していた活性化を MP 次数で分割。100B パラメータ・MP=16 で活性化メモリが 33 GB → 2 GB。Pa+cpu では CPU オフロードでほぼゼロ化。 - **CB(定数サイズバッファ)**: all-reduce 等の一時バッファをモデルサイズ非依存の定数サイズに固定。 - **MD(メモリ断片化解消)**: 活性化チェックポイントと勾配の連続領域予約・コピーにより、「空きメモリ 30% 超でも OOM」という断片化問題を解消。 ## 横断的知見 - **ZeRO は「DP のメモリ効率を MP レベルに引き上げる」が、逆に MP の必要性を弱める**: ZeRO が普及した後も MP は「極大モデルでの活性化削減」と「DP のみでは最小バッチサイズが過大になる場合の収束確保」という 2 つの残余価値を持つ。([[@2020__SC__ZeRO Memory Optimizations Toward Training Trillion Parameter Models]] §1) - **ZeRO の Stage と実装の通信プリミティブは論文の説明とずれることがある**: [[DeepSpeed]] の ZeRO Stage 2 は、論文上の説明と異なり既定では reduce-scatter ではなく all-reduce で勾配集約しており、`reduce_scatter=true` だけでなく `use_multi_rank_bucket_allreduce=false` を設定する必要があると後続研究([[@2025__PMBS__Pretraining LLMs at Scale - Tuning Strategies and Performance Portability]])が指摘する。「論文上の通信量」と「実際の実装の通信プリミティブ」のずれは、他フレームワーク(FSDP・Megatron・GPT-NeoX)でも起きる可能性がある。 - **ZeRO-DP の複製がデータ並列複製を「耐障害の冗長コピー」として再利用できる**: [[FlashRecovery]] の研究([[@2025__arXiv__FlashRecovery - Fast and Low-Cost Recovery from Failures for Large-Scale Training of LLMs]])は、ZeRO/FSDP 構成でデータ並列度 N のとき各デバイスに N−1 個のモデル状態の複製が存在することを利用し、障害時に同一 DP グループから集合通信で復元する設計を提案する。ZeRO の分割設計が「メモリ効率と耐障害性」という 2 つの価値を同時に提供できることを示す。 - **ZeRO stage・batch・gradient accumulation の組み合わせが実効スループットを大きく左右する**: [[@2025__PMBS__Pretraining LLMs at Scale - Tuning Strategies and Performance Portability]] では、同じ ZeRO-DP でも Stage 1 の batch_size 128 が OOM になり Stage 2 が最良になるなど、ハイパーパラメータ探索の重要性が示された。並列化戦略は「モデルをどう分割するか」だけでなく「その分割がどの通信プリミティブで実行されるか」まで含めて実測する必要がある。 ## 未解決の問い - Stage 3 の通信量 1.5 倍増加がボトルネックになる構成はどのような環境か。帯域・レイテンシ・モデルサイズの閾値は。 - ZeRO-DP のデータ並列複製と耐障害の複製冗長化の兼ね合い: 効率最適化(大きな Stage 3 で Nd を小さくしたい)と耐障害性(Nd を大きくしたい)が衝突する場合の最適化はどうなるか。 - FSDP(PyTorch 組み込み)と DeepSpeed ZeRO の実装差異が実性能に及ぼす影響の全体像。 - ZeRO-Offload(CPU/NVMe オフロード)や ZeRO-Infinity など、原論文の延長線上にある後続手法との統合設計の限界。 ## 関連 - 実装: [[DeepSpeed]] / FSDP(PyTorch) - 人物: [[Samyam Rajbhandari]] / [[Yuxiong He]] - 親概念: [[並列化戦略]] / [[LLM分散学習]] - 関連概念: [[集合通信]] / [[チェックポイント]] / [[耐障害LLM訓練]] - 関連 MOC: [[分散深層学習 - MOC]] ## 出典 - [[@2020__SC__ZeRO Memory Optimizations Toward Training Trillion Parameter Models]] — ZeRO の原論文。SC 2020。ZeRO-DP の 3 段階設計と ZeRO-R の詳細、400 V100 GPU での 100B 訓練実証。 - [[@2025__PMBS__Pretraining LLMs at Scale - Tuning Strategies and Performance Portability]] — ZeRO Stage・batch size・通信プリミティブの実測チューニングと DeepSpeed 実装の通信プリミティブ差異を報告。 - [[@2025__arXiv__FlashRecovery - Fast and Low-Cost Recovery from Failures for Large-Scale Training of LLMs]] — ZeRO/FSDP のデータ並列複製を耐障害冗長コピーとして再解釈し、チェックポイント不要な復旧手法を提案。