ゼロから作るDeep Learning ❻ ―LLM編

9章学習挑戦編

WebBotでは、StoryBotよりもさらに大きなモデルとデータを扱います。モデルが大きくなると、GPUメモリの制約が深刻な問題になります。1つのGPUに収まらなければ、複数のGPUを使う必要も出てきます。本章では「挑戦編」として、メモリ効率の改善、分散学習、より効率的なオプティマイザなど、大規模モデルの学習に必要な実践的なテクニックを紹介します。

9.1　勾配蓄積

学習時のGPUメモリ使用量は、バッチサイズに大きく依存します。バッチサイズが大きいほど、各層の中間出力（forwardの計算結果で、backward時に必要）が増え、メモリを圧迫します。かといってバッチサイズを小さくすると、学習の安定性や収束速度が低下します。

勾配蓄積（Gradient accumulation）は、この問題を解決するテクニックです。小さなバッチで複数回forward/backwardを行い、勾配を蓄積してからまとめてパラメータを更新します。これにより、実際のバッチサイズは小さくても、大きなバッチで学習したのと同等の効果が得られます。たとえば、バッチサイズ8で4回蓄積すれば、バッチサイズ32で学習したのと同じです。実装は以下のようになります。

accumulation_steps = 4
optimizer.zero_grad()

for i, (inputs, targets) in enumerate(dataloader):
    logits = model(inputs)
    loss = F.cross_entropy(logits, targets) / accumulation_steps  # 損失をスケーリング
    loss.backward()  # 勾配を蓄積

    if