AI・機械学習
中級

メモリ効率化訓練(メモリコウリツカクンレン)

限られたGPUメモリで大規模モデルを訓練するための総合的な最適化戦略。勾配チェックポインティング・Mixed Precision・ZeRO・Offloadingなど複数手法を組み合わせ、メモリ使用量を最小化する。

0 回閲覧
0 いいね

メモリ効率化訓練とは

メモリ効率化訓練(Memory-Efficient Training)は、GPUメモリの制約の中で可能な限り大きなモデルを効率的に訓練するための技術群の総称である。LLMの規模が年々拡大する一方で、単一GPUのメモリ容量(A100: 80GB、H100: 80GB、H200: 141GB)は限られているため、メモリ最適化は大規模モデル訓練の必須技術となっている。

GPUメモリの内訳は主に4つに分かれる:

メモリ区分内容7Bモデル(FP32)割合
パラメータモデルの重み28GB25%
勾配パラメータと同サイズ28GB25%
オプティマイザ状態Adam: パラメータの2倍56GB50%
活性値バッチサイズ・シーケンス長に依存可変可変

FP32の7BモデルをAdamで訓練すると、パラメータ+勾配+オプティマイザだけで112GBが必要。活性値を加えると単一A100では到底収まらない。メモリ効率化訓練はこの各区分を最適化する。

主要な最適化手法

1. Mixed Precision Training(混合精度訓練)

FP32の代わりにFP16またはBF16で計算することでメモリを半減させる。マスターウェイトはFP32で保持し、順伝播・逆伝播はFP16/BF16で実行する。Loss Scalingにより精度低下を防ぐ。

  • FP16: 広くサポートされるが、ダイナミックレンジが狭くオーバーフローリスクあり
  • BF16: A100以降のGPUでサポート。FP32と同じ指数部(8bit)を持ちオーバーフロー耐性が高い
  • メモリ削減効果: パラメータ+勾配+活性値が約半減

2. ZeRO(Zero Redundancy Optimizer)

DeepSpeedが実装するメモリ最適化手法。データ並列の各GPUが保持する冗長データを排除する。

ZeRO Stage分割対象メモリ削減(8GPU)
Stage 1オプティマイザ状態4倍
Stage 2+ 勾配8倍
Stage 3+ パラメータ線形スケーリング

ZeRO Stage 3では、各GPUがパラメータの1/N(N=GPU数)のみを保持し、必要時に他GPUから収集(All-Gather)する。通信コストは増加するが、メモリはGPU数に比例して削減される。

3. 勾配チェックポインティング

活性値メモリを削減する手法。チェックポイント間の活性値を逆伝播時に再計算する。メモリをO(N)からO(√N)に削減し、計算コストは20〜33%増加する。

4. CPU/NVMe Offloading

GPUメモリに収まらないデータをCPUメモリやNVMe SSDに退避する。DeepSpeed ZeRO-Offload、ZeRO-Infinity が代表的実装。

  • CPU Offload: オプティマイザ状態と勾配をCPUメモリに移動。PCIe帯域がボトルネック
  • NVMe Offload: さらにNVMe SSDを活用。テラバイト級のメモリ空間を利用可能

5. 勾配累積(Gradient Accumulation)

バッチサイズを分割して複数のマイクロバッチに分け、各マイクロバッチの勾配を累積してからパラメータ更新する。実効バッチサイズを維持しつつ、1ステップあたりのメモリ使用量を削減する。

手法の組み合わせ実例

7Bモデルのフルファインチューニング(A100 80GB×8)

  • ZeRO Stage 3: パラメータ・勾配・オプティマイザを8分割
  • BF16 Mixed Precision: 数値精度をBF16に
  • 勾配チェックポインティング: 全Transformerブロックに適用
  • 勾配累積: 4ステップ累積でバッチサイズ128相当
  • FlashAttention-2: Attentionメモリ最適化
  • 結果: 各GPU約45GB使用、80GB以内に収まる

70Bモデルのファインチューニング(A100 80GB×8)

上記に加えて:

  • LoRA(rank=16): 訓練パラメータを0.1%に削減
  • CPU Offload: オプティマイザ状態をCPUに退避
  • 結果: 各GPU約65GB使用

FAQ

Q1: どの最適化手法を最初に試すべきですか?

まずMixed Precision(BF16)と勾配チェックポインティングを有効化する。この2つだけでメモリを60〜70%削減でき、実装コストも低い。それでもメモリが不足する場合にZeROやOffloadingを追加する。

Q2: メモリ効率化は訓練速度を低下させますか?

手法による。Mixed Precisionは演算を高速化するためむしろ速度向上する。勾配チェックポインティングは20〜33%の速度低下。ZeRO Stage 1-2はオーバーヘッド5%以下。CPU Offloadingは通信ボトルネックで30〜50%の速度低下が起きうる。

Q3: LoRAとフルファインチューニングでメモリ使用量はどう異なりますか?

LoRA(rank=16)は訓練対象パラメータが全体の0.1%程度になるため、勾配とオプティマイザ状態のメモリが99%削減される。7BモデルでLoRA使用時は単一A100 80GBで訓練可能。フルファインチューニングでは8GPU以上が必要。