AI・機械学習
中級

Memory Offloading(メモリ オフローディング)

深層学習の学習時に GPU メモリに収まりきらないテンソル(オプティマイザ状態、勾配、パラメータ、活性化値)の一部を CPU メモリや NVMe SSD に退避させる手法。DeepSpeed ZeRO-Offload や PyTorch FSDP の CPU Offload で実装され、限られた GPU でも大規模モデルの学習を可能にする。

0 回閲覧
0 いいね

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 の帯域が利用可能。

速度への影響

オフロードは速度とメモリのトレードオフです。

オフロード構成メモリ削減速度低下
オプティマイザのみ → CPU50-60%10-20%
オプティマイザ + 勾配 → CPU70-80%20-40%
全コンポーネント → CPU80-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 など)を試すのが推奨です。量子化はメモリ削減と速度向上を同時に実現できるため、オフロードの速度ペナルティを回避できます。量子化でも不十分な場合にオフロードを追加するのがベストプラクティスです。