Gradient Accumulation

勾配累積は、マイクロバッチと呼ばれる小さなバッチを複数回、順伝播と誤差逆伝播に通し、勾配を加算してからモデルの重みを更新するトレーニング手法です。この更新は、すべてのマイクロバッチを合わせた大きさのバッチを使った1回のステップを近似します。その間、GPUに保持されるのは一度に1つのマイクロバッチ分のアクティベーションだけです。画像解像度を下げたりモデルを小さくしたりせずに、メモリに制約のあるデバイスでレシピが指定するバッチサイズを使ってトレーニングできます。実効バッチサイズは、マイクロバッチサイズ×累積ステップ数×デバイス数です。したがって、4枚の画像からなるマイクロバッチを4台のGPUで8ステップ累積すると、バッチサイズ128として動作します。

勾配蓄積の仕組み#

PyTorchやほとんどの他のフレームワークでは、backward()を呼び出すたびに、パラメーターに保存済みの勾配を置き換えず、新たな勾配を加算します。勾配累積はこの仕組みをそのまま利用します。各マイクロバッチで順伝播と逆伝播を実行し、所定の数に達したらオプティマイザーを1回実行してから、勾配をクリアし、次のウィンドウを開始します。

ほとんどの損失関数はバッチ内のサンプル数で平均するため、マイクロバッチの損失をステップ数で割り、加算した勾配を平均にします。このスケーリングを行えば、8個ずつのマイクロバッチ4つは、バッチサイズ32の1つのバッチと同じ勾配を生成します。

import torch
from torch import nn

torch.manual_seed(0)
model = nn.Linear(16, 1)
x, y = torch.randn(32, 16), torch.randn(32, 1)
loss_fn = nn.MSELoss()

# Reference: one batch of 32 samples
loss_fn(model(x), y).backward()
batch_grad = model.weight.grad.clone()
model.zero_grad()

# Four micro-batches of 8: backward() adds into .grad, and dividing each loss by 4 turns the sum into a mean
steps = 4
for x_micro, y_micro in zip(x.chunk(steps), y.chunk(steps)):
    (loss_fn(model(x_micro), y_micro) / steps).backward()

print("Matches the full-batch gradient:", torch.allclose(model.weight.grad, batch_grad, atol=1e-6))
print(f"Largest difference: {(model.weight.grad - batch_grad).abs().max():.1e}")

最大差1.6e-07でMatches the full-batch gradient: Trueが出力されます。これは同じ数値を異なる順序で加算したことによる浮動小数点の丸め誤差です。

大きなバッチと勾配累積が異なる点#

独立したサンプルを平均する損失関数では両者は一致しますが、よくある3つのケースでは一致しません。

• バッチ正規化: マイクロバッチごとに平均と分散を計算するため、4枚の画像から得られる統計量は128枚の場合よりノイズが大きくなります。
• 可変長シーケンス: 言語モデルが各マイクロバッチのトークン数で損失を平均すると、トークン数の少ないマイクロバッチも、長いマイクロバッチと同じ重みになります。
• 勾配クリッピング: オプティマイザーステップの直前に、累積した勾配に対して1回だけ適用します。各マイクロバッチでクリッピングすると、結果が変わります。

Hugging FaceとUnslothは、大規模言語モデル（LLM）のファインチューニングにおけるトークン数の不一致を2024年10月に修正しました。累積ウィンドウ全体の合計トークン数で、損失の合計を割る方法です。勾配累積による高速化はありません。同じ数のサンプルが、依然として1つずつマイクロバッチでネットワークを通過し、小さなマイクロバッチではGPUの利用効率も低くなります。分散トレーニングでは、PyTorchのDistributedDataParallelにあるno_sync()コンテキストを使うと、最後のマイクロバッチ以外では勾配のall-reduceをスキップでき、通信を削減できます。実効バッチサイズを大きくする場合は、通常、学習率を高くする必要があります。Facebookの研究者は、256台のGPUで画像8,192枚のバッチを使い、バッチサイズに比例して学習率をスケーリングし、段階的にウォームアップすることで、ImageNetでResNet-50を1時間でトレーニングしました。

関連手法#

• 勾配チェックポイント: アクティベーションを保存せずに、誤差逆伝播中に再計算します。追加の計算コストと引き換えに、各マイクロバッチで使用するメモリを節約します。
• 混合精度: アクティベーションを16ビット形式で保存することで、より大きなマイクロバッチを使えるようにします。損失スケーラーは、最後のマイクロバッチの後に勾配を一度だけスケーリング解除します。
• パイプライン並列化: 各バッチをマイクロバッチに分割し、別々のGPU上にあるモデルの各段階に通します。オプティマイザーステップを1回実行する前に、それらの勾配を累積します。

Ultralytics YOLO26のトレーニングでは、既定で勾配を累積します。トレーニングモードでは、トレーナーはround(nbs / batch)バッチごとにオプティマイザーを実行します。そのため、既定のbatch=16と公称バッチサイズnbs=64では、更新ごとに4バッチが処理され、batch=8では8バッチが処理されます。ウォームアップ中は1から段階的に増え、weight decayはbatch × accumulation / nbsでスケーリングされます。そのため、64枚の画像を想定した既定のハイパーパラメーターは、一度に数枚しか保持できないGPUでも適用できます。