Gradient Checkpointing

グラディエントチェックポイントは、順伝播の中間アクティベーションの一部だけを保存し、残りを誤差逆伝播中に再計算することでメモリを節約するトレーニング手法です。アクティベーションチェックポイントとも呼ばれ、追加の計算と引き換えにピークメモリ使用量を抑えます。この手法は名前に反して、パラメーターの勾配やモデルファイルではなくアクティベーションをチェックポイント化するため、アクティベーションテンソルが原因でニューラルネットワークを利用可能なGPUメモリに収められない場合に特に有効です。

グラディエントチェックポイントの仕組み#

標準的な順伝播では、ニューラルネットワークがアクティベーションと呼ばれる中間テンソルを計算します。PyTorchのautogradの仕組みで説明されているように、自動微分システムは、後で勾配を計算するために必要なアクティベーションを保持します。ディープネットワーク、大きなバッチ、高解像度画像、長い入力シーケンスでは、保存されたテンソルが大量のメモリを消費する場合があります。

グラディエントチェックポイントでは、ネットワークをセグメントに分割します。

1. 順伝播では、選択したセグメントの入力または境界のアクティベーションを保存します。
2. それらのセグメント内にある他の中間アクティベーションは破棄します。
3. 逆伝播では、各チェックポイント対象セグメントで順伝播を再実行し、失われた値を再構築します。
4. 再構築されたアクティベーションは、すぐに勾配の計算に使用します。

PyTorchのアクティベーションチェックポイントAPIでは、torch.utils.checkpointを通じてこの動作を利用できます。同等の概念として、JAXのグラディエントチェックポイントと再マテリアライゼーションや、TensorFlowのテープチェックポイントがあります。

チェックポイントの配置によってトレードオフが決まります。チェックポイント対象の領域を増やすほど、一般にメモリは節約できますが、繰り返される演算も増えます。ネットワーク全体を再計算するより、アクティベーションの多いブロックの周囲に選択的に配置するほうが、バランスを取りやすい場合があります。

トレードオフと関連手法#

グラディエントチェックポイントを使っても、通常はモデルパラメーター、勾配、オプティマイザーの状態は変わりません。主にアクティベーションメモリを削減する手法であり、一部の順伝播計算を2回実行するため、トレーニングは遅くなります。正確な結果は、アーキテクチャ、チェックポイント境界、バッチ形状、ハードウェアによって異なります。

これは、いくつかの関連手法とは異なります。

• グラディエント累積は、モデルの重みを更新する前に複数のマイクロバッチを処理することで、すべてのサンプルを同時に読み込まずに、より大きな実効バッチを実現します。一方、グラディエントチェックポイントは、各マイクロバッチで保持するアクティベーションを削減します。
• 混合精度は、選択した演算に低精度のデータ型を使用します。PyTorchの自動混合精度ワークフローは、メモリを削減して対応する演算を高速化できる一方、チェックポイントは意図的に計算量を増やします。
• バッチサイズの削減は、まとめて処理するサンプル数を減らしてメモリを節約します。チェックポイントを使うと、より大きなバッチや入力解像度を実用可能にできる場合があります。
• トレーニングチェックポイントは、復旧や後の推論に備えて重みとオプティマイザーの状態を保存します。PyTorchのモデルチェックポイントガイドでは、この永続化の仕組みを説明しています。これはアクティベーションの再計算とは関係ありません。

チェックポイント対象のコードは、元の順伝播と再計算時の順伝播で機能的に一貫している必要があります。チェックポイント対象の領域内に変更可能な状態、デバイス間転送、制御されていない乱数があると、エラーや不正な勾配が発生する場合があります。デコーダー型Transformerでは、キャッシュされた推論状態が順伝播グラフの再構築と競合する場合があるため、チェックポイントを使ったトレーニング中にkey-valueキャッシュを無効にする必要があることもあります。

PyTorchの例#

次の例では、1回のトレーニングステップ中にメモリ負荷の高いブロックをチェックポイント化します。

import torch
from torch import nn
from torch.utils.checkpoint import checkpoint

torch.manual_seed(0)

block = nn.Sequential(
    nn.Linear(1024, 4096),
    nn.ReLU(),
    nn.Linear(4096, 1024),
)
optimizer = torch.optim.AdamW(block.parameters())
inputs = torch.randn(8, 1024, requires_grad=True)
targets = torch.zeros_like(inputs)

optimizer.zero_grad(set_to_none=True)
outputs = checkpoint(block, inputs, use_reentrant=False)
loss = nn.functional.mse_loss(outputs, targets)
loss.backward()
optimizer.step()

ブロックの入力と必要な境界情報だけを保持し、内部のアクティベーションはloss.backward()中に再構築します。実際のプロジェクトでは、PyTorchのGPUメモリ最大割り当て量などの指標を使い、チェックポイントを使用した場合と使用しない場合の実行を比較する必要があります。

実際のアプリケーション#

• 高解像度コンピュータービジョン: 医療画像のセグメンテーション、航空画像の検出、産業検査では、特徴マップがモデルの重みより多くのメモリを消費する大きな画像を使ってトレーニングする場合があります。バックボーンの一部のステージにチェックポイントを適用すると、入力を大幅に縮小する代わりに、画像解像度を維持したり、バッチあたりのサンプル数を増やしたりできます。

• 長いシーケンスを使ったTransformerのトレーニング: シーケンス長や層数が増えると、アクティベーションの保存量は急速に増加します。Transformerブロックを再計算することで、同じアクセラレーター上でも、より長いコンテキストや大きなマイクロバッチを扱えるようになります。NVIDIAのアクティベーション再計算ガイドでは、Transformer層全体の再計算と選択的な再計算を紹介しています。

実践的な指針#

プロファイリングによって、パラメーターやオプティマイザーの状態ではなく、アクティベーションがメモリの大部分を占めると分かった場合に、グラディエントチェックポイントを使用してください。まず大きな反復ブロックから始め、対象範囲を広げる前に、ピークメモリ、反復時間、検証時の挙動をベンチマークしてください。

文書化されたUltralytics YOLOのトレーニング設定では、まず自動混合精度、画像サイズ、物理バッチサイズをメモリ制御に利用できます。Ultralytics AutoBatchのリファレンスでは、利用可能なGPUメモリに基づいてバッチサイズを自動選択する方法を説明しています。ローカルハードウェアでメモリが不足する場合は、Ultralytics Platformのクラウドトレーニングで、管理されたトレーニング実行向けに設定可能なクラウドGPUを利用できます。カスタムPyTorchアーキテクチャでアクティベーションメモリをより細かく管理する必要がある場合は、グラディエントチェックポイントをこれらの制御と組み合わせられます。