> [!abstract] 概要 > 大規模深層学習モデルは精度向上に有効だが、数十億から兆規模のパラメータを訓練することは困難である。データ並列・モデル並列などの既存手法は、限られたデバイスメモリへの適合、計算効率、通信効率、開発効率の面で根本的な制約を持つ。本論文では、新手法 ZeRO(Zero Redundancy Optimizer)を提案する。ZeRO はメモリを最適化し、訓練速度を大幅に向上させながら、効率的に訓練できるモデルサイズを拡大する。ZeRO はデータ並列・モデル並列訓練におけるメモリ冗長性を排除しつつ、低い通信量と高い計算粒度を維持することで、デバイス数に比例したモデルサイズのスケーリングを持続的な高効率で実現する。メモリ要件と通信量の分析より、ZeRO は現行ハードウェアで 1 兆パラメータを超えるスケーリングの可能性を持つ。ZeRO を実装・評価した結果、400 GPU 上で 100B パラメータ超の大規模モデルを超線形スピードアップで訓練し、15 ペタフロップスのスループットを達成した。これは SOTA 比でモデルサイズ 8 倍・性能 10 倍以上の向上に相当する。ユーザビリティの面では、ZeRO はモデル並列を必要とせずに最大 13B パラメータ(Megatron GPT 8.3B や T5 11B より大きい)のモデルを訓練でき、データサイエンティストが大規模モデルを自由に実験できる環境を提供する。ZeRO の成果として、当時世界最大の言語モデル(17B パラメータ、記録更新の精度)が誕生した。 ## 論文情報 | 項目 | 内容 | |---|---| | 論文名 | ZeRO: Memory Optimizations Toward Training Trillion Parameter Models | | 著者 | Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Yuxiong He (Microsoft) | | 掲載 | arXiv:1910.02054 → SC 2020 | | 公開日 | 2020-05-13(v3) | | URL | https://arxiv.org/abs/1910.02054 | | 実装 | [[DeepSpeed]](Microsoft OSS 訓練最適化ライブラリ) | ## 概要 大規模言語モデル(LLM)の訓練において、デバイスメモリが最大のボトルネックとなっている。本論文は、メモリ消費を「モデル状態」(オプティマイザ状態・勾配・パラメータ)と「残余状態」(活性化・一時バッファ・断片化メモリ)の 2 種類に分類し、それぞれに異なるアプローチで最適化する ZeRO システムを提案する。 ZeRO はデータ並列の通信効率とモデル並列のメモリ効率を同時に達成する点に新規性があり、SC 2020 で発表後 [[DeepSpeed]] の中核として実装された。 ## 問題設定 ### 既存手法の限界 - **データ並列(DP)**: 各デバイスにモデル状態を複製するため、1.4B パラメータ以上で 32 GB GPU でもメモリ不足になる。通信効率と計算粒度は高いが、メモリ効率は低い。 - **モデル並列(MP)**: 層内でパラメータを分割するため、ノード内(高帯域 NVSwitch)でのみ効率的。ノード間では帯域が 300 GB/秒から 12.5 GB/秒(InfiniBand EDR)に急落し、40B パラメータ超で効率が激減する。実測で Megatron-LM の 40B モデル(DGX-2 × 2 ノード)で GPU 1 つあたり約 5 テラフロップス(ハードウェア峰値の 5% 未満)。 - **パイプライン並列(PP)**: GPipe はバブル隠蔽のためにバッチサイズ依存性があり活性化メモリが大きい。PipeDream は古くなったパラメータの複数コピーを保持してメモリ非効率。 ### メモリ消費の定量化 混合精度(fp16/fp32)訓練と Adam オプティマイザの組み合わせでは、Ψ パラメータのモデルに対して以下のメモリが必要: - fp16 パラメータ: 2Ψ バイト - fp16 勾配: 2Ψ バイト - fp32 パラメータ複製(オプティマイザ用): 4Ψ バイト - fp32 モメンタム(Adam): 4Ψ バイト - fp32 分散(Adam): 4Ψ バイト - **合計: 16Ψ バイト**(K=12 の場合、K はオプティマイザ状態のメモリ乗数) GPT-2(1.5B パラメータ)でも 24 GB 以上を要し、単一 GPU(32 GB)での訓練が事実上不可能。 ## 提案手法 ### ZeRO-DP: モデル状態のゼロ冗長最適化 3 段階の分割でメモリ冗長を排除する。Nd を DP 次数とする。 **Stage 1 — オプティマイザ状態分割(P_os)** - オプティマイザ状態を Nd 均等分割し、各プロセスは自身の分割分(1/Nd)のみ更新・保持 - 訓練ステップ末尾に all-gather でパラメータを集約 - メモリ削減: 4Ψ + KΨ → 4Ψ + KΨ/Nd → **約 4 倍**削減(Nd が大きい場合) - 通信量: DP と同じ 2Ψ **Stage 2 — 勾配分割(P_os+g)** - 各プロセスは自身が担当するパラメータ分割に対応する勾配のみ保持 - 後退伝播中にバケット化戦略で reduce-scatter を実行 - メモリ削減: **約 8 倍** - 通信量: DP と同じ 2Ψ(reduce-scatter Ψ + all-gather Ψ) **Stage 3 — パラメータ分割(P_os+g+p)** - 各プロセスは自身の分割分のパラメータのみ保持 - 順伝播・後退伝播時にパラメータをブロードキャストで受け取り、使用後に破棄 - メモリ削減: **Nd 倍**(1T パラメータを 1024 GPU で訓練可能) - 通信量: 1.5 倍(3Ψ、基準 DP の 1.5 倍) | DP 次数(Nd) | 標準 DP | P_os | P_os+g | P_os+g+p | |---|---|---|---|---| | 1 | 120 GB | 120 GB | 120 GB | 120 GB | | 64 | 120 GB | 31.4 GB | 16.6 GB | **1.88 GB** | | 1024 | 120 GB | 30.1 GB | 15.1 GB | **0.12 GB** | *(7.5B パラメータモデルの例、32 GB V100 GPU)* ### ZeRO-R: 残余メモリの最適化 **Pa — 分割活性化チェックポイント** - MP 設計では活性化がすべての MP GPU に複製されるが、ZeRO-R はこれを分割して保持 - 後退伝播時に all-gather で再構成 - メモリ削減: MP 次数に比例(MP=16 で 100B モデルの活性化を 33 GB → 2 GB) - Pa+cpu: 分割チェックポイントを CPU にオフロードし活性化メモリをほぼゼロ化 **CB — 定数サイズバッファ** - 大規模 all-reduce 等の一時バッファが「モデルサイズに比例」して肥大化する問題を解消 - 「性能効率を保ちつつ十分な大きさ」の定数サイズ融合バッファを使用 - 例: 3B パラメータモデルの fp32 バッファが 12 GB → 定数上限に抑制 **MD — メモリ断片化解消** - 順伝播の活性化チェックポイント(長寿命)と再計算活性化(短寿命)の混在が断片化を引き起こす - 断片化で「空きメモリが 30% 超あるのに OOM」という問題が実際に発生 - 活性化チェックポイントと勾配のためにあらかじめ確保した連続メモリ領域へコピーすることで解消 ### ZeRO と MP の組み合わせ ZeRO-DP は MP と組み合わせ可能で、各デバイスのメモリを最大 Nd × Nm 倍削減できる(Nm は MP 次数)。1024 GPU で 16-way MP + 64-way DP により 1 兆パラメータモデルの訓練が理論上可能。ただし ZeRO がメモリ問題を解決した今、MP の主な残余価値は「極大モデルでの活性化削減」と「DP のみでは最小バッチサイズが小さすぎる場合の収束確保」である。 ## 新規性 1. **DP のメモリ効率を MP レベルまで引き上げながら通信量を DP レベルに維持**する設計は先行研究にない。既存手法はメモリ効率と通信効率のどちらかを犠牲にしていた。 2. **モデル変更不要**のインターフェース設計: `torch.nn.Module` を ZeRO でラップするだけで利用可能。MP や PP のようなモデル改造が不要。 3. **ZeRO-R の 3 要素(Pa/CB/MD)の統合**: 活性化・バッファ・断片化という残余メモリの 3 源泉をすべて体系的に解決。 4. **超線形スケーラビリティ**: DP 次数増加でモデル状態のメモリが削減される → バッチサイズを大きく取れる → アリスメティック強度上昇によってスループットが超線形に増加。 ## 実験設定 - **ハードウェア**: 400 V100 GPU(25 DGX-2 ノード)、ノード間通信帯域 800 Gbps - **実装**: ZeRO-100B(P_os+g + ZeRO-R)を PyTorch で実装。`torch.nn.Module` 互換インターフェース - **ベースライン**: MP なし = PyTorch DDP、MP あり = Megatron-LM(2019年9月 OSS 版) - **モデル**: GPT-2 ライクな Transformer。パラメータ数は隠れ次元とレイヤー数で変化(1.5B〜170B) ## 実験結果 ### モデルサイズ - ZeRO-100B + MP: **170B パラメータ**まで効率的に実行(Megatron-LM 単独では 40B 超で効率低下) - **8 倍以上**のモデルサイズ向上 ### スループット - 100B パラメータモデルを 400 GPU で **38 テラフロップス/GPU**、**15 ペタフロップス**の集計性能 - SOTA 比で同モデルサイズで **10 倍以上**の速度向上 - 64〜400 GPU のレンジで**超線形スピードアップ**を観察(GPU 数を倍にするとスループットが 2 倍超に増加) ### 民主化効果 - MP・PP なしで最大 **13B パラメータ**を訓練可能(T5 11B より大きい) - PyTorch DDP は 1.4B パラメータで OOM になるが、ZeRO は 9 倍以上のモデルを処理 - MP 不要なためノード内高帯域インターコネクト(NVLINK/NVSwitch)が不要になる ### Turing-NLG - ZeRO-100B で **17B パラメータ**の Turing-NLG を訓練(2020年5月時点で世界最大) - Webtext-103 ペルプレキシティ **10.21**(Megatron-LM 8.3B の SOTA を更新) - 41.4 テラフロップス/GPU の持続スループット ## 考察 ZeRO は「データ並列の全複製」と「モデル並列の全分割」という二項対立を解消する。鍵は「**すべての状態を常時全プロセスに保持する必要はない**」という洞察であり、オプティマイザ状態は訓練ステップ末尾のみ必要、勾配は後退伝播直後のみ必要、パラメータは順伝播/後退伝播の該当レイヤーのみ必要という時間的性質を活かして動的通信スケジュールを設計した。 Stage 1〜2 では通信量を一切増やさずにメモリを削減できるため、既存 DP ベースのワークロードへの移行コストが低い。Stage 3 は通信量 1.5 倍増加を伴うが、それでも MP のノード間通信に比べれば大幅に効率的。 [[DeepSpeed]] への組み込みにより、ZeRO は後続研究・大規模モデル訓練のデファクトスタンダードとなった。 ## 強み - **モデル変更不要**という圧倒的なユーザビリティ。MP/PP は全て不要 - Stage 1〜2 で通信量変化なしのメモリ削減という優れたトレードオフ - ZeRO-DP と ZeRO-R を組み合わせることで、モデル状態と残余状態の両面を網羅 - [[DeepSpeed]] として OSS 化され、実証された実務的インパクト - 超線形スケーラビリティという付加価値(GPU 増加でバッチサイズも増えてスループットが加速) ## 弱点・課題 - Stage 3(パラメータ分割)は通信量 1.5 倍増加を招く。通信帯域が細い環境では不利 - 実装は ZeRO-100B(Stage 1〜2 + ZeRO-R)のサブセットにとどまり、Stage 3 は SC 2020 時点で未実装(DeepSpeed の後バージョンで対応) - 1 兆パラメータは理論的に可能でも、実計算量から訓練時間が 1 年超になり計算資源が課題 - Pa+cpu(活性化を CPU にオフロード)は帯域制約のある PCI-E 経由でコストがかかる - Stage 3 はパラメータをブロードキャストで受け取るため、通信実装の複雑性が増す - 大規模バッチサイズによる収束への悪影響(critical batch size 超の場合)は ZeRO の範囲外