Memory Offloading(メモリ オフローディング)
深層学習の学習時に GPU メモリに収まりきらないテンソル(オプティマイザ状態、勾配、パラメータ、活性化値)の一部を CPU メモリや NVMe SSD に退避させる手法。DeepSpeed ZeRO-Offload や PyTorch FSDP の CPU Offload で実装され、限られた GPU でも大規模モデルの学習を可能にする。
Memory Offloading とは
Memory Offloading は、GPU の VRAM に収まらないデータを CPU の RAM や NVMe SSD などの低速・大容量ストレージに退避させることで、単一 GPU でもより大きなモデルを学習できるようにする手法です。GPU の VRAM が 24GB(RTX 4090)しかなくても、CPU メモリが 128GB あれば、そこにオプティマイザ状態を退避させることで 30B パラメータ規模のモデルを扱えます。
オフロード対象の優先順位
学習時に GPU メモリを消費する主要コンポーネントとそのオフロード優先度を整理します。
| コンポーネント | メモリ消費(7B FP16) | アクセス頻度 | オフロード効果 |
|---|---|---|---|
| オプティマイザ状態 | 約 56GB(Adam) | 更新時のみ | 最も効果的 |
| 勾配 | 約 14GB | 逆伝播時 | 効果的 |
| モデルパラメータ | 約 14GB | 毎ステップ | 中程度 |
| 活性化値 | 約 60GB | 順伝播・逆伝播 | チェックポイント併用推奨 |
Adam オプティマイザはパラメータごとに momentum と variance の 2 つの FP32 状態を持つため、モデルパラメータの 4 倍のメモリを消費します。このオプティマイザ状態を CPU にオフロードするだけで GPU メモリの大部分を解放できます。
DeepSpeed ZeRO-Offload
Microsoft の DeepSpeed フレームワークが提供する ZeRO-Offload は、Memory Offloading の最も成熟した実装です。
ZeRO-Offload(Stage 2): オプティマイザ状態と勾配を CPU にオフロード。パラメータ更新は CPU 上で FP32 で実行し、更新後のパラメータを GPU に転送。
ZeRO-Infinity(Stage 3): パラメータ自体も CPU/NVMe にオフロード。理論上は無制限のモデルサイズに対応。ただし通信オーバーヘッドが大きい。
NVMe Offload: CPU メモリにも収まらない場合、NVMe SSD にさらにオフロード。PCIe Gen4 NVMe で約 7GB/s の帯域が利用可能。
速度への影響
オフロードは速度とメモリのトレードオフです。
| オフロード構成 | メモリ削減 | 速度低下 |
|---|---|---|
| オプティマイザのみ → CPU | 50-60% | 10-20% |
| オプティマイザ + 勾配 → CPU | 70-80% | 20-40% |
| 全コンポーネント → CPU | 80-90% | 40-60% |
| CPU + NVMe 併用 | 90%+ | 60-80% |
速度低下の主因は GPU-CPU 間の PCIe 帯域(Gen4 x16 で約 32GB/s)がボトルネックになることです。最新の PCIe Gen5 では約 64GB/s に倍増し、オフロードの速度ペナルティが改善されます。
CPU 側の要件
オフロードを効果的に活用するには CPU 側のリソースも重要です。
- RAM 容量: オフロード対象のテンソルサイズ以上が必要。7B モデルのオプティマイザオフロードなら最低 64GB、13B なら 128GB 推奨
- CPU コア数: パラメータ更新を CPU で行う場合、8 コア以上推奨。更新の並列化でオーバーヘッドを軽減
- メモリ帯域: DDR5-4800 で約 38.4GB/s。大量のテンソルの読み書きが発生するため高帯域メモリが有利
- NVMe 速度: NVMe オフロード時は Gen4 以上の高速 SSD を推奨。Gen3 SSD では帯域不足で速度低下が顕著
個人環境での活用例
RTX 4090(24GB VRAM)+ 128GB DDR5 RAM の自作 PC で、DeepSpeed ZeRO-Offload を使用した実例です。
- 7B モデルフルファインチューニング: オプティマイザ + 勾配 CPU オフロードで VRAM 使用量を約 18GB に抑制。速度は純 GPU の約 65%
- 13B モデル LoRA: オプティマイザ CPU オフロードで VRAM 約 20GB。速度は純 GPU の約 85%
- 30B モデル推論: パラメータの一部を CPU にオフロードして VRAM 22GB で動作。生成速度は約 5 tokens/s
FAQ
Q1: Apple Silicon Mac でもメモリオフロードは使えますか?
A1: Apple Silicon は統合メモリ(Unified Memory)を採用しているため、GPU と CPU が同じメモリプールを共有します。従来のオフロード概念とは異なりますが、PyTorch MPS バックエンドで Metal Performance Shaders を活用した学習が可能です。M4 Max の 128GB 統合メモリでは 70B モデルの推論も報告されています。
Q2: 複数 GPU とオフロードの組み合わせは有効ですか?
A2: はい。DeepSpeed ZeRO-3 + Offload の組み合わせで、各 GPU のメモリ使用量を最小化しつつ CPU にもオフロードできます。2 × RTX 4090 + 128GB RAM で 70B モデルの LoRA ファインチューニングが可能になります。
Q3: オフロードと量子化はどちらを先に試すべきですか?
A3: まず量子化(QLoRA など)を試すのが推奨です。量子化はメモリ削減と速度向上を同時に実現できるため、オフロードの速度ペナルティを回避できます。量子化でも不十分な場合にオフロードを追加するのがベストプラクティスです。