Microsoft Research が開発した分散学習最適化技術。Zero Redundancy Optimizer の略で、データ並列学習時のメモリ冗長性を排除し、数千GPUへのスケーリングを可能にする。Stage 1〜3の段階的なメモリ削減とCPU/NVMeオフロードを提供する。
DeepSpeed ZeRO（Zero Redundancy Optimizer）は、Microsoft Researchが2020年に発表した分散学習メモリ最適化技術である。従来のデータ並列（DDP）では各GPUにモデルパラメータ・勾配・オプティマイザ状態の完全なコピーを保持するため、メモリ効率が極めて低かった。ZeROはこれらの状態をGPU間で分割（シャーディング）し、メモリ冗長性をゼロにすることで、同一GPU数でより大規模なモデルの学習を可能にする。
Mixed Precision学習（FP16/BF16）環境で、1パラメータあたりのメモリ消費を分析すると、モデルパラメータ（2B）、勾配（2B）、オプティマイザ状態（Adam: FP32パラメータ4B + FP32一次モーメント4B + FP32二次モーメント4B = 12B）で合計16バイトとなる。10Bモデルでは160GBのメモリが必要で、8基のA100 80GBでもデータ並列では収まらない。
ZeROは3つのステージで段階的にメモリを削減する。各ステージの特性を理解し、モデルサイズとGPU数に応じて最適なステージを選択することが重要である。
| ステージ | 分割対象 | メモリ削減率 | 通信量（vs DDP） | 推奨用途 |
|---|---|---|---|---|
| Stage 1（ZeRO-1） | オプティマイザ状態のみ | 最大4倍 | 同等 | 〜30Bモデル、少数GPU |
| Stage 2（ZeRO-2） | オプティマイザ状態 + 勾配 | 最大8倍 | 同等 | 〜70Bモデル、中規模クラスタ |
| Stage 3（ZeRO-3） | オプティマイザ状態 + 勾配 + パラメータ | 最大N倍（GPU数比例） | 1.5倍 | 100B+モデル、大規模クラスタ |
各GPUがオプティマイザ状態（Adamの場合12B/パラメータ）の1/Nのみを保持する。Forward/Backwardは通常のDDPと同様に実行し、AllReduce後に各GPUが担当するパラメータのオプティマイザ更新のみ行う。通信量はDDPと完全に同等で、オーバーヘッドなくメモリを最大4倍削減できる。
Stage 1に加え、勾配もGPU間で分割する。Backward時にReduceScatterで勾配を各GPUの担当パラメータに集約し、不要な勾配メモリを即座に解放する。通信パターンがAllReduceからReduceScatter + AllGatherに変わるが、総通信量は同等である。
パラメータ自体もGPU間で分割する最も積極的なステージ。Forward/Backward時に必要なパラメータをAllGatherで一時的に収集し、計算後に解放する。GPU数Nに比例してメモリが削減されるため、理論上は任意サイズのモデルを学習可能。ただし、AllGather通信が追加されるため、通信量はDDPの約1.5倍に増加する。
ZeRO-Infinityは、GPUメモリの制約をCPUメモリとNVMe SSDに拡張する技術である。Stage 3と組み合わせて使用し、使用していないパラメータ・勾配・オプティマイザ状態をCPUメモリまたはNVMe SSDにオフロードする。
| オフロード先 | 帯域幅 | 容量（典型） | 適用シナリオ |
|---|---|---|---|
| GPUメモリのみ | 3.35 TB/s（HBM3e） | 80-192 GB/GPU | 標準的な学習 |
| CPU DRAM | 200 GB/s（DDR5） | 512 GB-2 TB/ノード | 大規模モデル・少数GPU |
| NVMe SSD | 7-14 GB/s（Gen4/5） | 2-8 TB/ノード | 超大規模モデル・メモリ極限 |
CPUオフロードでは、パラメータのプリフェッチ（先読み）と非同期データ転送により、GPU計算と通信をオーバーラップさせてオーバーヘッドを最小化する。NVMeオフロードではさらに帯域が制限されるため、効率的なI/Oパイプラインが不可欠である。
DeepSpeedの設定はJSONファイルで管理し、HuggingFace Transformersとの統合が容易である。
{
"bf16": {"enabled": true},
"zero_optimization": {
"stage": 3,
"offload_optimizer": {"device": "cpu", "pin_memory": true},
"offload_param": {"device": "cpu", "pin_memory": true},
"overlap_comm": true,
"contiguous_gradients": true,
"sub_group_size": 1e9,
"reduce_bucket_size": "auto",
"stage3_prefetch_bucket_size": "auto",
"stage3_param_persistence_threshold": "auto",
"stage3_max_live_parameters": 1e9,
"stage3_max_reuse_distance": 1e9
},
"gradient_accumulation_steps": 4,
"train_micro_batch_size_per_gpu": 2,
"wall_clock_breakdown": false
}
HuggingFace Trainer統合では --deepspeed ds_config.json を指定するだけで、自動的にZeROが有効化される。
DeepSpeedエコシステムには、Mixture-of-Experts（MoE）モデル向けの並列化（DeepSpeed-MoE）と、長シーケンス学習向けのシーケンス並列（DeepSpeed-Ulysses）も含まれる。
DeepSpeed-Ulyssesは、Attention計算のシーケンス次元をGPU間で分割し、QKV変換後にAll-to-All通信でHead次元の分割に切り替える。これにより、コンテキスト長100万トークン以上の長シーケンス学習が可能となる。
モデルがStage 2でGPUメモリに収まるなら、Stage 2を推奨する。Stage 3は追加のAllGather通信が発生するため、Stage 2より5〜15%スループットが低下する。70B以上のモデルや、GPUメモリが限られる環境（A100 40GB等）ではStage 3が必要となる。
ZeRO-Offloadはオプティマイザ計算自体をCPUで実行する「CPU Optimizer」アプローチも提供する。FSDPのCPUオフロードはパラメータの一時退避が主で、計算はGPU上で行う。ZeRO-OffloadはCPU計算とGPU計算のパイプライン化が高度に最適化されており、少数GPU環境でのコスト効率が高い。
Megatron-DeepSpeedとして公式に統合されている。Megatron-LMのテンソル並列・パイプライン並列とDeepSpeed ZeROのデータ並列を組み合わせた3D並列が実現できる。NVIDIA NeMoフレームワークもこの組み合わせを標準サポートしている。