AI・機械学習
上級

勾配累積(コウバイルイセキ)

複数のマイクロバッチにわたるフォワード・バックワードパスの勾配を加算し、蓄積完了後に一括してパラメータ更新を行う手法。GPU メモリに収まらない大きなバッチサイズを擬似的に実現でき、単一 GPU でも大バッチ学習の恩恵を得られる。LLM のファインチューニングからフロンティアモデルの事前学習まで、あらゆるスケールで利用される基盤技術である。

0 回閲覧
0 いいね

勾配累積とは

勾配累積(Gradient Accumulation)は、深層学習におけるメモリ効率化の基本手法であり、1 回のパラメータ更新に必要な勾配を複数回のフォワード・バックワードパスに分割して計算・蓄積する技術である。数式で表すと、累積ステップ数を K、各ステップの勾配を g_k とした場合、パラメータ更新に使用する勾配は G = (1/K) × Σ(k=1 to K) g_k となり、これはバッチサイズ K 倍での一括計算と数学的に等価である。

LLM の文脈では、GPU メモリの大部分がモデルパラメータとオプティマイザ状態で消費されるため、データに割り当てられるメモリが極めて限られる。たとえば 7B パラメータモデルを AdamW で学習する場合、パラメータ(FP16: 14GB)+ オプティマイザ状態(FP32: 84GB)= 98GB となり、A100 80GB の単一 GPU にはモデルすら載らない。ZeRO Stage 2 でオプティマイザ状態を 8 GPU に分散しても、1 GPU あたりの空きメモリは 20〜30GB 程度で、シーケンス長 4,096 のマイクロバッチサイズは 2〜4 が限界となる。ここで勾配累積 16 ステップを適用すれば、実効バッチサイズ 32〜64 × 8 GPU = 256〜512 を実現できる。

勾配累積の数学的等価性

勾配累積が大バッチと数学的に等価であることを示す。損失関数 L の全バッチ B に対する勾配は以下の通り:

∇L(B) = (1/|B|) × Σ(x∈B) ∇l(x)

これを K 個のサブバッチ B_1, ..., B_K に分割すると:

∇L(B) = (1/K) × Σ(k=1 to K) [(1/|B_k|) × Σ(x∈B_k) ∇l(x)]

各サブバッチの勾配 g_k = (1/|B_k|) × Σ(x∈B_k) ∇l(x) を計算し、K 回分を平均すれば全バッチの勾配と一致する。

ただし、Batch Normalization(BN)を使用する場合はサブバッチごとに正規化統計量が異なるため厳密な等価性は崩れる。LLM では BN ではなく Layer Normalization(LN)を使用するため、勾配累積の等価性は完全に保たれる。

実装の詳細

PyTorch での勾配累積の基本実装パターンは以下の構造に従う:

  1. 各マイクロバッチでフォワードパスを実行し損失を計算
  2. 損失を累積ステップ数 K で割ってからバックワードパスを実行(勾配のスケーリング)
  3. K ステップ分の勾配が蓄積されたらオプティマイザのステップを実行
  4. 勾配をゼロクリアして次のサイクルを開始

重要な実装上の注意点:

  • 損失のスケーリング: loss = loss / K を忘れると勾配が K 倍に膨らみ、学習が発散する。最も多いバグの一つ
  • 勾配クリッピングのタイミング: 勾配クリッピング(max_grad_norm)は累積完了後、optimizer.step() の直前に適用する。各マイクロバッチ後にクリッピングすると実効的なクリップ閾値が変わる
  • 学習率スケジューラのステップ: scheduler.step() はパラメータ更新ごと(K マイクロバッチに 1 回)に呼ぶ。各マイクロバッチで呼ぶと学習率が K 倍速く減衰する
  • 混合精度との連携: AMP(Automatic Mixed Precision)使用時は GradScaler のスケーリングも累積に対応させる必要がある。PyTorch 2.x 以降では GradScaler が自動的に処理

メモリ消費の比較

勾配累積によるメモリ節約効果を定量的に示す。7B パラメータモデル、シーケンス長 2,048 の場合:

設定マイクロバッチサイズ累積ステップ実効バッチサイズ活性化メモリ総メモリ
大バッチ直接32132約 48GB約 146GB
勾配累積(中)4832約 6GB約 104GB
勾配累積(小)13232約 1.5GB約 99.5GB
累積 + チェックポイント4832約 1.2GB約 99.2GB

活性化メモリはマイクロバッチサイズに比例するため、勾配累積により大幅に削減できる。さらに活性化チェックポインティング(再計算法)を併用すれば、活性化メモリを大幅に圧縮可能。

勾配累積のオーバーヘッド

勾配累積は数学的に等価だが、実行時間にはオーバーヘッドが生じる:

計算オーバーヘッド:フォワード・バックワードパスの計算量自体は同一だが、K 回のカーネル起動オーバーヘッドが加算される。GPU の CUDA カーネル起動は数マイクロ秒だが、小さなマイクロバッチでは演算時間自体も短いため相対的に無視できない。一般に 3〜8% の実行時間増加。

メモリアクセスオーバーヘッド:勾配テンソルへの繰り返し加算は追加のメモリ読み書きを発生させる。ただし勾配テンソルは L1/L2 キャッシュに収まることが多く、実質的な影響は小さい。

通信の非重畳:DDP(Distributed Data Parallel)では通常バックワードパスと勾配の All-Reduce 通信を重畳して行うが、勾配累積中は通信を累積完了まで遅延させる必要がある。no_sync() コンテキストマネージャで中間ステップの通信を抑制し、最終ステップでのみ All-Reduce を実行する。

オーバーヘッド種別影響度対策
カーネル起動3〜8%CUDA Graph で軽減可
メモリ帯域1〜3%勾配 in-place 加算
通信遅延(DDP)5〜15%no_sync() で中間通信抑制
合計5〜20%累積ステップ数が少ないほど小さい

Hugging Face Transformers での利用

Hugging Face の TrainingArguments では gradient_accumulation_steps パラメータで直接指定できる。per_device_train_batch_size=4 と gradient_accumulation_steps=8 を 4 GPU で使用すれば、グローバルバッチサイズは 4 × 8 × 4 = 128 となる。Trainer クラスが損失のスケーリング、勾配クリッピングのタイミング、スケジューラのステップ管理、DDP の no_sync() をすべて自動処理するため、実装ミスのリスクが大幅に軽減される。

2026 年時点では Trainer の gradient_accumulation_steps は整数のみ対応だが、学習途中での動的変更は TrainerCallback 内で args.gradient_accumulation_steps を書き換えることで可能。バッチサイズウォームアップとの組み合わせで利用される。

混合精度・量子化との連携

勾配累積は混合精度学習(AMP)や量子化学習と組み合わせることで、さらなるメモリ節約を実現する:

  • FP16 AMP: フォワードパスを FP16、勾配累積を FP32 で行うことで、活性化メモリを半減しつつ累積精度を維持
  • BF16: 指数部が FP32 と同じ 8 ビットのため、勾配累積の数値安定性が FP16 より高い。H100 以降の GPU で推奨
  • FP8 学習: H100 の FP8 Tensor Core を活用。勾配累積は FP32 で行い、フォワード・バックワードの活性化のみ FP8 にキャスト
  • QLoRA: 4-bit 量子化されたベースモデルに対し LoRA アダプタの勾配を FP16/BF16 で累積。メモリ効率が極めて高く、24GB GPU(RTX 4090)で 70B モデルのファインチューニングが可能

FAQ

Q: 勾配累積のステップ数に上限はありますか?

A: 理論的な上限はないが、実用上は 32〜128 ステップが一般的な範囲である。ステップ数が極端に大きい(256+)場合、1 回のパラメータ更新に要する壁時計時間が長くなり、学習率スケジューラの解像度が低下する。また、FP16 で勾配を累積する場合、多数回の加算で浮動小数点の丸め誤差が蓄積される可能性がある。BF16 または FP32 での累積を推奨。DeepSpeed では fp32_accumulation オプションで FP32 累積を強制できる。

Q: 勾配累積はバッチ正規化と相性が悪いですか?

A: はい。バッチ正規化はバッチ内のサンプルで平均・分散を計算するため、小さなマイクロバッチではこれらの統計量が不安定になり、大バッチでの直接計算と結果が異なる。しかし LLM では Layer Normalization を使用するため、この問題は発生しない。Layer Normalization は各サンプルの特徴次元内で正規化するため、バッチサイズに依存せず、勾配累積との完全な互換性がある。

Q: 勾配累積中にデータの順序はシャッフルされますか?

A: 各エポック開始時のデータシャッフルは通常通り行われ、マイクロバッチはシャッフル後のデータから順に取り出される。勾配累積中のマイクロバッチ間で追加のシャッフルは行われないが、エポック単位のシャッフルにより十分なランダム性が確保される。ただし、分散学習では各 GPU が異なるデータサブセットを受け取るよう DistributedSampler でシャッフルシードを同期する必要がある。

Q: 勾配累積はオンライン学習(ストリーミングデータ)でも使えますか?

A: 可能だが、注意点がある。オンライン学習ではデータがストリームで到着するため、累積中のマイクロバッチ間でデータ分布が変動する可能性がある。特に時系列データや推薦システムのリアルタイム学習では、累積ステップ数を小さく保つ(4〜8)ことで分布ドリフトの影響を最小化する。LLM の継続学習(continual pre-training)でも同様の考慮が必要。