Gradient Accumulation(グラディエントアキュムレーション)
複数のミニバッチの勾配を蓄積してからパラメータ更新を行う手法。GPU メモリに収まらない大きな実効バッチサイズを、メモリ制約内で実現する分散訓練の基盤技術。
Gradient Accumulation(勾配蓄積)とは
Gradient Accumulationは、1回のパラメータ更新(optimizer step)に複数のミニバッチの勾配を蓄積する手法です。GPUメモリの制約でバッチサイズを大きくできない場合に、実効的な大バッチ訓練を可能にします。LLMの訓練ではGPUメモリが常にボトルネックであり、Gradient Accumulationは事実上すべてのLLM訓練パイプラインで使用されています。
なぜ大バッチサイズが必要か
LLMの訓練では、バッチサイズが大きいほど勾配の推定が安定し、訓練効率が向上します。GPT-3ではバッチサイズ3.2M tokensが使用されましたが、これを1つのGPUでそのまま処理することは不可能です。
| モデル | 実効バッチサイズ | GPU メモリ | 物理バッチ | 蓄積ステップ |
|---|---|---|---|---|
| GPT-3 175B | 3.2M tokens | 80GB A100 | 1-2 sequences | 数百〜数千 |
| LLaMA 7B | 4M tokens | 80GB A100 | 8-16 sequences | 32-64 |
| BERT Base | 256 sequences | 16GB V100 | 32 sequences | 8 |
動作原理
通常の訓練ループでは、各ミニバッチで以下を実行します。
- Forward pass(順伝播)
- Loss計算
- Backward pass(逆伝播・勾配計算)
- Optimizer step(パラメータ更新)
- 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 | 蓄積ステップ |
|---|---|---|---|
| 256 | 1 | 8 | 32 |
| 256 | 4 | 8 | 8 |
| 256 | 8 | 8 | 4 |
| 1024 | 8 | 16 | 8 |
蓄積ステップが大きいほど更新頻度が下がり、訓練時間が延びます。GPUメモリが許す限り物理バッチサイズを最大化し、蓄積ステップを最小化するのが効率的です。
FAQ
Q1: Gradient Accumulationは訓練速度に影響しますか?
計算量自体は変わりませんが、optimizer stepの頻度が減るため、学習率スケジュールのステップ単位が変わります。蓄積ステップを考慮してスケジュールを調整する必要があります。
Q2: 蓄積ステップ数が大きすぎるとどうなりますか?
パラメータ更新の頻度が下がりすぎると、収束が遅くなります。また、勾配の蓄積中にモデルの重みが更新されないため、古い重みでの勾配計算が続くことで若干の最適化効率の低下が生じます。
Q3: Mixed Precision(FP16/BF16)との併用で注意点はありますか?
GradScalerを使用するAMP環境では、Gradient AccumulationとGradScalerの相互作用に注意が必要です。蓄積中のlossスケーリングは各ミニバッチで行い、unscaleはoptimizer step時にのみ実行します。