AI・機械学習
中級

Gradient Accumulation(グラディエントアキュムレーション)

複数のミニバッチの勾配を蓄積してからパラメータ更新を行う手法。GPU メモリに収まらない大きな実効バッチサイズを、メモリ制約内で実現する分散訓練の基盤技術。

0 回閲覧
0 いいね

Gradient Accumulation(勾配蓄積)とは

Gradient Accumulationは、1回のパラメータ更新(optimizer step)に複数のミニバッチの勾配を蓄積する手法です。GPUメモリの制約でバッチサイズを大きくできない場合に、実効的な大バッチ訓練を可能にします。LLMの訓練ではGPUメモリが常にボトルネックであり、Gradient Accumulationは事実上すべてのLLM訓練パイプラインで使用されています。

なぜ大バッチサイズが必要か

LLMの訓練では、バッチサイズが大きいほど勾配の推定が安定し、訓練効率が向上します。GPT-3ではバッチサイズ3.2M tokensが使用されましたが、これを1つのGPUでそのまま処理することは不可能です。

モデル実効バッチサイズGPU メモリ物理バッチ蓄積ステップ
GPT-3 175B3.2M tokens80GB A1001-2 sequences数百〜数千
LLaMA 7B4M tokens80GB A1008-16 sequences32-64
BERT Base256 sequences16GB V10032 sequences8

動作原理

通常の訓練ループでは、各ミニバッチで以下を実行します。

  1. Forward pass(順伝播)
  2. Loss計算
  3. Backward pass(逆伝播・勾配計算)
  4. Optimizer step(パラメータ更新)
  5. Gradient zero(勾配リセット)

Gradient Accumulationでは、ステップ4-5をN回に1回だけ実行します。N回分の勾配が蓄積されてから初めてパラメータを更新するため、実効バッチサイズはN倍になります。

実装のポイント

Loss の正規化

蓄積時はlossをaccumulation_stepsで割る必要があります。これを忘れると、実効的な学習率がaccumulation_steps倍になり、訓練が不安定になります。

loss = model(batch).loss / accumulation_steps
loss.backward()

if (step + 1) % accumulation_steps == 0:
    optimizer.step()
    optimizer.zero_grad()

バッチ正規化との非互換

Batch Normalizationの統計量(mean/variance)はミニバッチ単位で計算されるため、Gradient Accumulationで実効バッチサイズを大きくしてもBNの統計量には反映されません。LLMではLayer Normalizationが使用されるため、この問題は発生しません。

Data Parallel との組み合わせ

分散訓練(DDP / FSDP)とGradient Accumulationは組み合わせて使用するのが一般的です。

手法効果メモリ影響
Gradient Accumulation実効バッチサイズ増大変化なし
Data Parallel (DDP)GPU間でバッチ分割GPU数分に分散
FSDPモデル+勾配をGPU間で分割大幅削減
3手法併用最大バッチサイズ実現最大効率

例: 8 GPU × 物理バッチ4 × 蓄積ステップ8 = 実効バッチサイズ256

Accumulation Steps の設定指針

目標バッチサイズGPU数物理バッチ/GPU蓄積ステップ
2561832
256488
256884
10248168

蓄積ステップが大きいほど更新頻度が下がり、訓練時間が延びます。GPUメモリが許す限り物理バッチサイズを最大化し、蓄積ステップを最小化するのが効率的です。

FAQ

Q1: Gradient Accumulationは訓練速度に影響しますか?

計算量自体は変わりませんが、optimizer stepの頻度が減るため、学習率スケジュールのステップ単位が変わります。蓄積ステップを考慮してスケジュールを調整する必要があります。

Q2: 蓄積ステップ数が大きすぎるとどうなりますか?

パラメータ更新の頻度が下がりすぎると、収束が遅くなります。また、勾配の蓄積中にモデルの重みが更新されないため、古い重みでの勾配計算が続くことで若干の最適化効率の低下が生じます。

Q3: Mixed Precision(FP16/BF16)との併用で注意点はありますか?

GradScalerを使用するAMP環境では、Gradient AccumulationとGradScalerの相互作用に注意が必要です。蓄積中のlossスケーリングは各ミニバッチで行い、unscaleはoptimizer step時にのみ実行します。