AI・機械学習
上級

CPUオフローディング(訓練)(シーピーユーオフローディング)

LLM訓練時にGPUメモリに収まらないデータ(オプティマイザ状態・勾配・パラメータ)をCPUメモリやNVMe SSDに退避し、必要時にGPUに転送して計算する手法。DeepSpeed ZeRO-OffloadやZeRO-Infinityが代表的実装。

0 回閲覧
0 いいね

CPUオフローディング(訓練)とは

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メモリの内訳(Llama 3 8B, FP16, Adam)

メモリ区分サイズGPUに残すCPUにオフロード
パラメータ(FP16)16GB○ZeRO-3で可
勾配(FP16)16GB○ZeRO-Offloadで可
Adam状態(FP32 m+v)64GB×◎(最優先対象)
活性値可変○△(帯域ボトルネック)

Adam系オプティマイザはパラメータの4倍のメモリを消費する(1次モーメント+2次モーメント、各FP32)。これが最もオフロード効果の高い対象であり、ZeRO-Offloadの主要ターゲットとなっている。

メモリ削減効果の実例

構成オフロードなしZeRO-OffloadZeRO-Infinity
1×V100 32GB最大1.4Bモデル最大10Bモデル−
1×A100 80GB最大7Bモデル最大30Bモデル最大70Bモデル
8×A100 80GB最大70Bモデル最大200Bモデル最大1Tモデル

DeepSpeed ZeRO-Offload

アーキテクチャ

ZeRO-Offloadは、GPUとCPU間のデータフローを最適化するために、以下の設計原則を採用する:

  1. 計算はGPUに集中: 順伝播と逆伝播の計算コストが高い処理は全てGPU上で実行
  2. 状態管理はCPUに委譲: Adamのパラメータ更新(FP32精度)はCPU上で実行
  3. 通信は非同期: GPU→CPUの勾配転送とCPU→GPUのパラメータ転送をパイプライン化

通信ボトルネック

PCIe Gen4 x16の帯域は理論値32GB/s(実測25〜28GB/s)。8Bモデルの勾配(16GB)転送に約0.6秒かかる。順伝播・逆伝播の計算時間(バッチサイズ依存で0.5〜2秒)と重なるようにパイプライン化することで、通信オーバーヘッドを隠蔽する。

PCIe Gen5ではこの帯域が2倍(64GB/s)に向上し、オフローディングの性能が大幅に改善される。

ZeRO-Infinity(NVMe Offload)

CPUメモリの限界を超える

CPUメモリ(512GB〜2TB)でも足りない超大規模モデルでは、NVMe SSD(数TB〜数十TB)をオフロード先として活用する。NVMe SSDの帯域は3〜7GB/s(PCIeの1/4〜1/10)と遅いが、容量は事実上無制限。

階層的メモリ管理

ZeRO-Infinityは3層のメモリ階層を管理する:

階層デバイス帯域容量用途
L1GPU VRAM2TB/s80-141GB計算中のテンソル
L2CPU DRAM25-100GB/s512GB-2TBオプティマイザ状態
L3NVMe SSD3-7GB/s数TBパラメータの一部

必要なデータを計算タイミングに合わせてL3→L2→L1にプリフェッチし、不要になったデータを逆方向に退避する。このパイプライン制御がZeRO-Infinityの核心。

実用上の注意点

スループットへの影響

CPUオフローディングは通信帯域がボトルネックとなるため、GPU onlyの訓練と比較してスループットが30〜50%低下する。NVMe Offloadではさらに低下する。ただし「訓練できないモデルが訓練可能になる」という点では唯一の選択肢となる場合がある。

CPUメモリと帯域の設計

オフローディング使用時は十分なCPUメモリと高速なメモリ帯域が必要。DDR5-4800の8チャネル構成では約300GB/sの帯域が得られ、PCIeボトルネックを回避できる。CPUメモリが不足するとスワップが発生し、性能が壊滅的に低下する。

FAQ

Q1: CPUオフローディングとモデル並列はどちらを選ぶべきですか?

複数のGPUが利用可能な場合はモデル並列(テンソル並列+パイプライン並列)の方がスループットが高い。CPUオフローディングは「GPUの枚数を増やせないがGPUメモリが足りない」場合の救済策として位置づけられる。両者の併用も可能で、DeepSpeedでは3D並列 + ZeRO-Offloadの構成がサポートされている。

Q2: CPUでのオプティマイザ計算はボトルネックになりませんか?

CPUでのAdam更新は並列化可能で、32コア以上のサーバーCPUであればGPUの計算時間と重なるようにパイプライン化できる。ただしCPUコア数が少ない環境(デスクトップPCの8コア等)では、オプティマイザ計算がボトルネックになり得る。

Q3: Apple Silicon MacでCPUオフローディングは使えますか?

Apple Silicon(M1/M2/M3/M4)はUnified Memory Architectureで、CPUとGPUが同一メモリ空間を共有するため、従来のPCIe転送が不要。MLXフレームワークでは自動的にUnified Memoryを活用するが、DeepSpeedのCPUオフローディングとは異なるアーキテクチャであり、直接的な互換性はない。