勾配チェックポイント

フォワードパスでは、バックワードパスで再利用するためにすべての中間アクティベーションをキャッシュするのが一般的です。ただし、アクティベーションはバッチサイズとシーケンス長に比例して増加します。Gradient checkpointingでは、一部のアクティベーションのみを保存し、残りは破棄します。これにより、バックワードパスでは、必要に応じて一部のアクティベーションをオンザフライで再計算する必要があります。

Normal training:
  Forward:   [L1]→[L2]→[L3]→[L4]   (save ALL activations)
  Backward:  ←uses cached activations everywhere

Gradient checkpointing:
  Forward:   [L1]→[L2]→[L3]→[L4]   (save only at checkpoints, discard the rest)
  Backward:  ←reaches L2, recomputes L2→L3 from scratch, uses it, discards it

バックワードパスで破棄されたアクティベーションを再計算するため、通常はトレーニングが遅くなりますが、チェックポイント化によりアクティベーションのメモリ使用量を削減できます。

gradient_checkpointing=True を設定して有効にします。

勾配蓄積と組み合わせて使用し、メモリ使用量をさらに削減できます。

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    gradient_checkpointing=True,
)

部分的なチェックポイント

Full gradient checkpointing は、チェックポイント可能なすべてのレイヤーを再計算します。実行にメモリに余裕がある場合は、チェックポイントするレイヤーを減らし、メモリ削減の一部を速度向上に振り分けることができます。

every_n_layers を gradient_checkpointing_enable() に渡すことで、チェックポイントの間隔を選択できます。

every_n_layers=2

Forward:   input  -> [L1] -> [L2] -> [L3] -> [L4] -> [L5] -> [L6]
                      CP     keep     CP     keep     CP     keep

Backward:  output <- [L6] <- [L5] <- [L4] <- [L3] <- [L2] <- [L1]
                     keep   rerun    keep   rerun    keep   rerun

every_n_layers=2 を指定すると、最初の層とその後の2層ごとにチェックポイントが設定されます。チェックポイントが設定された層は、フォワードパス中にアクティベーションを破棄し、バックワードパス中に再計算します。一方、その他の層はアクティベーションをメモリに保持します。

model.gradient_checkpointing_enable(every_n_layers=2)

デフォルトのevery_n_layers=1では、すべてのレイヤーがチェックポイント化されます。より大きな値を指定すると、最初のレイヤーがチェックポイント化され、その後、nレイヤーごとにチェックポイント化されます。その他のレイヤーのアクティベーションはメモリ内に保持されます。例えば、every_n_layers=2を指定すると、レイヤー 1、3、5 などがチェックポイント化されます。カウントされるのは、GradientCheckpointingLayerを継承するモジュールのみです。グラデーションチェックポイントに対応するその他のモジュールは、引き続き有効なままです。

Trainerで部分的な勾配チェックポイントを使用するには、gradient_checkpointing_kwargs内のevery_n_layersを設定します。

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"every_n_layers": 4},
)

保存されたアクティベーションのオフロード

Gradient checkpointing は、チェックポイント化された各レイヤーごとに 1 つのアクティベーションを GPU に保持します。シーケンス長が長い場合、これにより大量のメモリを消費する可能性があります（約 layers x sequence x hidden x bytes_per_element）。offload を設定すると、これらのアクティベーションをピン留めされたホストメモリに保持できます。その代わりとして、フォワードパスではデバイスからホストへのコピー、バックワードパスではホストからデバイスへのコピーが発生します。

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"offload": True},
)

両方のコピーは compute stream 上で実行されるため、メモリを消費する代わりにステップが遅くなります。実行がそれ以外の方法では収まらない場合に使用してください。高速化を目的として使用しないでください。

次のステップ

• GPU上のメモリ使用量を駆動している要因を理解するために、GPUメモリ使用量のドキュメントをお読みください。
• メモリをさらに削減し、トレーニングを高速化するために、低精度のデータ型を使用する方法については、Mixed precision trainingガイドを参照してください。
• カスタム融合カーネルによる学習の高速化方法については、Kernelsガイドを参照してください。