勾配の蓄積

大きなバッチサイズは大きなアクティベーションを生成し、GPU メモリを使い切ってしまいます。勾配蓄積（Gradient accumulation）を使用すると、勾配の計算を複数のミニバッチに分散させることで、より大きな実効バッチサイズでの学習が可能になります。

勾配は、オプティマイザーが重みを更新する前に、n個のミニバッチにわたって蓄積されます。例えば、デバイスごとのバッチサイズが8で、蓄積ステップ数が4の場合、実効バッチサイズは32になります。

Step 1: mini-batch 1 → forward → backward → grads = G₁
Step 2: mini-batch 2 → forward → backward → grads = G₁ + G₂
Step 3: mini-batch 3 → forward → backward → grads = G₁ + G₂ + G₃
Step 4: mini-batch 4 → forward → backward → grads = G₁ + G₂ + G₃ + G₄
        → optimizer.step()  ← same update as if batch_size × 4
        → zero_grad()

グラデーションの蓄積は、より大きなバッチがメモリに収まらない場合のみ使用してください。これは、真に大きなバッチでの学習と比較してスループットを向上させるものではありません。

per_device_train_batch_size に対して gradient_accumulation_steps 回にわたって勾配を累積します。

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    per_device_train_batch_size=8,
    gradient_accumulation_steps=4,
)

損失スケーリング

カスタム損失関数を使用する場合は、num_items_in_batchを含めることで、Trainerがすべてのミニバッチにわたる予測ターゲットの数で損失を割るようにします。これにより、gradient_accumulation_stepsによる固定のステップ数ではなく、トークン数で正規化されます。それ以外の場合、Trainerは損失をgradient_accumulation_stepsで割ります。

import torch.nn.functional as F

def compute_loss(outputs, labels, num_items_in_batch=None):
    logits = outputs["logits"]
    loss = F.cross_entropy(logits, labels, reduction="sum")
    return loss / num_items_in_batch

因果言語モデルの場合、num_items_in_batch は シフトされた ラベルを数えます。損失関数はラベルをシフトするため、位置 i での予測は位置 i + 1 のトークンを対象とし、これにより各シーケンスの位置 0 にはターゲットが存在しません。Trainer はこれらの位置を除外し、labels[..., 1:] に対して集計を行うため、分母は損失関数が使用する予測ターゲットの数と一致します。データコラテーターが shift_labels を直接提供する場合は、例えばパディングなしのコラテーターなど、Trainer はそのテンソルに対して集計を行います。マスクド言語モデルや分類などの他の損失タイプでは、ラベルテンソル全体を数えます。

Trainer をサブクラス化して compute_loss() をオーバーライドする場合、loss_is_scaled_for_ga を設定して、損失がどのように正規化されているかを指定します。False はバッチごとの平均値を意味し、Trainer が gradient_accumulation_steps で割ります。True は、すでに num_items_in_batch で割っている場合です。

class MyTrainer(Trainer):
    loss_is_scaled_for_ga = True

    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        ...
        return loss_sum / num_items_in_batch

次のステップ

GPU上のメモリ使用量を駆動している要因を理解するために、GPUメモリ使用量のドキュメントをお読みください。
アクティベーションをキャッシュする代わりに再計算することでアクティベーションメモリを削減する方法を学ぶには、Gradient checkpointingガイドを参照してください。
メモリを削減し、トレーニングを高速化するために低精度のデータ型を使用する方法については、Mixed precision trainingガイドを参照してください。
勾配の計算方法について詳しく知るには、Gradient Accumulation Fixのブログ記事をお読みください。

Update on GitHub

Transformers

勾配の蓄積

損失スケーリング

次のステップ