LLM訓練時にGPUメモリに収まらないデータ(オプティマイザ状態・勾配・パラメータ)をCPUメモリやNVMe SSDに退避し、必要時にGPUに転送して計算する手法。DeepSpeed ZeRO-OffloadやZeRO-Infinityが代表的実装。
CPUオフローディングは、GPUメモリの容量制限を超えて大規模モデルを訓練するための手法である。GPUメモリ(VRAM)は最大でも80〜141GB(H200)だが、CPUメモリ(DRAM)はサーバーで512GB〜2TBを搭載可能。この大容量のCPUメモリをGPUの拡張領域として活用し、モデルの訓練データの一部をCPUに退避(オフロード)する。
DeepSpeedが2021年に発表したZeRO-Offloadが先駆的な実装で、オプティマイザ状態と勾配計算の一部をCPUに移動することで、単一GPU上で10倍大きなモデルの訓練を可能にした。2021年にはZeRO-Infinityへと発展し、NVMe SSDへのオフロードにも対応してペタバイト級の仮想メモリ空間を実現している。
| メモリ区分 | サイズ | GPUに残す | CPUにオフロード |
|---|---|---|---|
| パラメータ(FP16) | 16GB | ○ | ZeRO-3で可 |
| 勾配(FP16) | 16GB | ○ | ZeRO-Offloadで可 |
| Adam状態(FP32 m+v) | 64GB | × | ◎(最優先対象) |
| 活性値 | 可変 | ○ | △(帯域ボトルネック) |
Adam系オプティマイザはパラメータの4倍のメモリを消費する(1次モーメント+2次モーメント、各FP32)。これが最もオフロード効果の高い対象であり、ZeRO-Offloadの主要ターゲットとなっている。
| 構成 | オフロードなし | ZeRO-Offload | ZeRO-Infinity |
|---|---|---|---|
| 1×V100 32GB | 最大1.4Bモデル | 最大10Bモデル | − |
| 1×A100 80GB | 最大7Bモデル | 最大30Bモデル | 最大70Bモデル |
| 8×A100 80GB | 最大70Bモデル | 最大200Bモデル | 最大1Tモデル |
ZeRO-Offloadは、GPUとCPU間のデータフローを最適化するために、以下の設計原則を採用する:
PCIe Gen4 x16の帯域は理論値32GB/s(実測25〜28GB/s)。8Bモデルの勾配(16GB)転送に約0.6秒かかる。順伝播・逆伝播の計算時間(バッチサイズ依存で0.5〜2秒)と重なるようにパイプライン化することで、通信オーバーヘッドを隠蔽する。
PCIe Gen5ではこの帯域が2倍(64GB/s)に向上し、オフローディングの性能が大幅に改善される。
CPUメモリ(512GB〜2TB)でも足りない超大規模モデルでは、NVMe SSD(数TB〜数十TB)をオフロード先として活用する。NVMe SSDの帯域は3〜7GB/s(PCIeの1/4〜1/10)と遅いが、容量は事実上無制限。
ZeRO-Infinityは3層のメモリ階層を管理する:
| 階層 | デバイス | 帯域 | 容量 | 用途 |
|---|---|---|---|---|
| L1 | GPU VRAM | 2TB/s | 80-141GB | 計算中のテンソル |
| L2 | CPU DRAM | 25-100GB/s | 512GB-2TB | オプティマイザ状態 |
| L3 | NVMe SSD | 3-7GB/s | 数TB | パラメータの一部 |
必要なデータを計算タイミングに合わせてL3→L2→L1にプリフェッチし、不要になったデータを逆方向に退避する。このパイプライン制御がZeRO-Infinityの核心。
CPUオフローディングは通信帯域がボトルネックとなるため、GPU onlyの訓練と比較してスループットが30〜50%低下する。NVMe Offloadではさらに低下する。ただし「訓練できないモデルが訓練可能になる」という点では唯一の選択肢となる場合がある。
オフローディング使用時は十分なCPUメモリと高速なメモリ帯域が必要。DDR5-4800の8チャネル構成では約300GB/sの帯域が得られ、PCIeボトルネックを回避できる。CPUメモリが不足するとスワップが発生し、性能が壊滅的に低下する。
複数のGPUが利用可能な場合はモデル並列(テンソル並列+パイプライン並列)の方がスループットが高い。CPUオフローディングは「GPUの枚数を増やせないがGPUメモリが足りない」場合の救済策として位置づけられる。両者の併用も可能で、DeepSpeedでは3D並列 + ZeRO-Offloadの構成がサポートされている。
CPUでのAdam更新は並列化可能で、32コア以上のサーバーCPUであればGPUの計算時間と重なるようにパイプライン化できる。ただしCPUコア数が少ない環境(デスクトップPCの8コア等)では、オプティマイザ計算がボトルネックになり得る。
Apple Silicon(M1/M2/M3/M4)はUnified Memory Architectureで、CPUとGPUが同一メモリ空間を共有するため、従来のPCIe転送が不要。MLXフレームワークでは自動的にUnified Memoryを活用するが、DeepSpeedのCPUオフローディングとは異なるアーキテクチャであり、直接的な互換性はない。