## Memo
- ![[Pasted image 20250428174159.png]]
## Memo with LLM
### 論文情報
- **タイトル**: ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
- **著者**: Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Yuxiong He (Microsoft)
- **カンファレンス/ジャーナル名**: arXiv preprint
- **発表年**: 2020年（arXiv投稿は2019年10月、更新は2020年5月）
### 論文概要
この論文は、大規模なディープラーニングモデルのトレーニングにおけるメモリ制約の問題に対処するために「Zero Redundancy Optimizer (ZeRO)」という新しい手法を提案しています。ZeROは、データ並列学習におけるメモリ冗長性を排除しながら、計算効率と通信量を維持することで、従来のアプローチと比較して最大8倍の大きさのモデル（100B以上のパラメータ）を10倍高速に学習できることを実証しています。
### 問題設定
#### 入力と出力
- **入力**: 数十億から1兆パラメータを持つ大規模ディープラーニングモデル
- **出力**: 効率的なトレーニングを可能にするメモリ最適化アプローチ
- **必要なデータ**: GPUクラスタのメモリ使用状況、通信ボリューム、計算効率に関するデータ
大規模モデルは精度向上に貢献しますが、単一デバイスのメモリには収まりきらないという問題があります。例えば、16ビット精度でのAdam最適化を使用する1.5Bパラメータモデルは、少なくとも24GBのメモリを必要としますが、これは単一のGPUで利用可能なメモリを超えています。
### 提案手法
ZeROは2つの主要な最適化セットから構成されています：
#### 1. ZeRO-DP (Data Parallelism)
モデル状態（オプティマイザー状態、勾配、パラメータ）のメモリ使用量を削減するための3段階の最適化：
1. **Pos (Optimizer State Partitioning)**: オプティマイザー状態を分割し、各データ並列プロセスが全体の1/Nd（Ndはデータ並列度）だけを更新・保存します。これにより最大4倍のメモリ削減が可能です。
2. **Pos+g (Gradient Partitioning)**: 勾配も分割し、各プロセスが担当パラメータの勾配のみを保持します。これにより最大8倍のメモリ削減が可能です。
3. **Pos+g+p (Parameter Partitioning)**: パラメータも分割し、Forward/Backward伝播時に必要なときだけ他プロセスから取得します。これにより、データ並列度Ndに比例したメモリ削減が可能になります（例：64GPUで64倍のメモリ削減）。
#### 2. ZeRO-R (Residual Memory Optimization)
残りのメモリ消費を最適化する手法：
1. **Pa (Partitioned Activation Checkpointing)**: 活性化チェックポイントを分割し、必要なときだけ全体を再構築します。これにより、モデル並列度に比例した活性化メモリの削減が可能です。
2. **CB (Constant Size Buffers)**: 計算効率とメモリ使用量のバランスを取るため、一時バッファサイズを固定します。
3. **MD (Memory Defragmentation)**: メモリフラグメンテーションを防ぐために、連続したメモリチャンクを事前に割り当て、活性化チェックポイントと勾配をその場でコピーします。
### 新規性
従来の手法との比較：
1. **データ並列処理(DP)との比較**: 標準的なDPはすべてのプロセスでモデル状態を複製するためメモリ効率が悪いのに対し、ZeROはモデル状態を分割しながらもDPの計算・通信効率を維持します。
2. **モデル並列処理(MP)との比較**: MPは計算の粒度を減らし通信オーバーヘッドを増加させる傾向があるのに対し、ZeROはより高い計算粒度と低い通信量を維持しながら同等以上のメモリ効率を実現します。
3. **パイプライン並列処理(PP)との比較**: PPはモデルを水平方向に分割しますが、バッチサイズ制約や実装の複雑さといった問題があります。ZeROはこれらの制約なしに同等または優れたメモリ効率を達成します。
ZeROの主要な新規性は、データ並列処理の計算効率と通信効率を維持しながら、モデル並列処理のようなメモリ効率を達成する点にあります。また、全てのモデル状態が常に必要なわけではないという洞察に基づき、動的通信スケジュールを活用している点も新規性があります。
### 実験設定
#### ハードウェア
- 400 NVIDIA V100 GPU（25 DGX-2ノード）
- ノード間通信帯域幅：800 Gbps
#### ベースライン
- モデル並列処理なしの実験：PyTorchのDistributed Data Parallel（DDP）
- モデル並列処理ありの実験：Megatron-LM（NVIDIAのオープンソース実装）
#### モデル
- GPT-2のようなトランスフォーマーベースのモデル
- 隠れ層の次元数とレイヤー数を変えることで、異なるパラメータ数のモデルを生成
#### 評価指標
- モデルサイズ（最大トレーニング可能なパラメータ数）
- スループット（GPU当たりのテラフロップス）
- スケーラビリティ（GPUの追加による性能向上）
### 実験結果
1. **モデルサイズ**: ZeRO-100B（Pos+gとZeRO-Rの実装）は400 GPUで最大170Bパラメータのモデルを効率的に学習可能で、これは従来のSOTA（Megatron-LM）の8倍以上です。
2. **スピード**: ZeRO-100Bは100Bパラメータモデルを400 NVIDIA V100 GPUクラスタで、GPU当たり38テラフロップス、合計15ペタフロップス以上のパフォーマンスで実行できます。これは同じモデルサイズに対して従来のSOTAと比較して10倍以上の高速化を実現しています。
3. **スケーラビリティ**: 64-400 GPUの範囲でスーパーリニアなスピードアップを観測しました。GPUの数を2倍にすると性能が2倍以上に向上しました。これはZeRO-DPがDP度を増やすとモデル状態のメモリフットプリントを減少させ、より大きなバッチサイズを効率的に処理できるようになるためです。
4. **大規模モデルトレーニングの民主化**: ZeRO-100Bは、モデル並列処理やパイプライン並列処理を必要とせずに、最大13Bパラメータのモデル（T5の11Bより大きい）をトレーニングできるため、データサイエンティストが複雑な並列処理を気にすることなく大規模モデルを実験できるようになりました。
5. **最新のSOTAモデル**: ZeROは17Bパラメータを持つ世界最大の言語モデルTuring-NLGの学習に利用され、記録的な精度を達成しました。
総合的に、ZeROは単に理論上だけでなく実装と評価において、大規模モデルトレーニングにおける重要なブレークスルーを示しています。アプローチを完全に実装すれば、現在のハードウェアで1兆パラメータのモデルをトレーニングすることも理論的には可能です。
## Abstract
大規模なディープラーニング・モデルは精度を大幅に向上させるが、数十億から数兆のパラメータをトレーニングすることは困難である。データやモデルの並列処理などの既存のソリューションでは、計算、通信、開発の効率を得ながら、これらのモデルを限られたデバイス・メモリに収めるには根本的な限界があります。私たちは、メモリを最適化し、学習速度を大幅に向上させながら、効率的に学習できるモデル・サイズを拡大する新しいソリューション、Zero Redundancy Optimizer（ZeRO）を開発します。ZeROは、低い通信量と高い計算粒度を維持しながら、データおよびモデル並列トレーニングにおけるメモリの冗長性を排除し、高い効率を維持しながら、デバイス数に比例してモデルサイズを拡張することを可能にする。メモリ要件と通信量に関する我々の分析は、次のことを示している： ZeROは、現在のハードウェアを使用して、1兆パラメータを超えるスケールを実現できる可能性があります。ZeROを実装して評価したところ、400個のGPUで100Bパラメーターを超える大規模モデルを超直線的なスピードアップで学習し、15ペタフロップスのスループットを達成しました。これは、モデルサイズが8倍、達成可能な性能が10倍向上したことを意味します。使い勝手の面では、ZeROは、科学者が適用しにくいモデル並列性を必要とすることなく、最大13Bパラメータ（例えば、Megatron GPT 8.3BやT5 11Bよりも大きい）の大規模モデルをトレーニングすることができる。最後になりましたが、研究者はZeROのシステムブレイクスルーを利用して、Turing-NLGを開発しました。