Gradient Checkpointing

勾配チェックポイントは、順伝播時の選択された中間アクティベーションのみを保存し、バックプロパゲーション中に残りを再計算する、メモリ節約型の学習手法です。アクティベーションチェックポイントとも呼ばれ、追加の計算処理と引き換えにピークメモリ使用量を削減します。その名前とは異なり、この手法はパラメータの勾配やモデルファイルではなくアクティベーションをチェックポイント化するため、アクティベーションテンソルによってニューラルネットワークを利用可能なGPUメモリに収められない場合に特に有効です。

勾配チェックポイントの仕組み#

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

勾配チェックポイントは、ネットワークをセグメントに分割します。

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

PyTorch activation checkpointing APIは、torch.utils.checkpointを通じてこの動作を公開しています。同様の概念として、JAX gradient checkpointing and rematerializationやTensorFlow tape checkpointingがあります。

チェックポイントの配置によってトレードオフが決まります。一般的に、より多くの領域をチェックポイント化すると多くのメモリが節約されますが、より多くの処理が繰り返されます。ネットワーク全体を再計算するよりも、アクティベーション集約型ブロックの周辺に選択的に配置する方が、より良いバランスを提供できます。

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

勾配チェックポイントは通常、モデルパラメータ、勾配、およびオプティマイザの状態を変更しません。主なターゲットはアクティベーションメモリであり、一部の順伝播計算が2回実行されるため、学習速度が低下します。正確な結果は、アーキテクチャ、チェックポイント境界、バッチ形状、およびハードウェアによって異なります。

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

• **勾配累積**は、モデルの重みを更新する前に複数のマイクロバッチを処理し、すべてのサンプルを同時に読み込むことなく、より大きな有効バッチを作成します。一方、勾配チェックポイントは、各マイクロバッチに対して保持されるアクティベーションを削減します。
• **混合精度**は、選択した操作に低精度のデータ型を使用します。PyTorch automatic mixed precision workflowはメモリを削減し、互換性のある操作を高速化できる一方、チェックポイントは意図的に計算を追加します。
• バッチサイズの削減は、処理するサンプル数を減らすことでメモリを削減します。チェックポイントを使用すると、より大きなバッチや入力解像度を維持できるようになる場合があります。
• 学習チェックポイントは、リカバリやその後の推論のために重みとオプティマイザの状態を保存します。PyTorch model checkpoint guideはこの永続化メカニズムについて説明しており、アクティベーションの再計算とは無関係です。

チェックポイント化されたコードは、元の順伝播実行と再計算された順伝播実行の間で機能的に一貫している必要があります。チェックポイント化された領域内のミュータブルな状態、デバイス転送、または制御されていないランダム性は、エラーや不正確な勾配を引き起こす可能性があります。デコーダー型のTransformerでは、キャッシュされた推論状態が順伝播グラフの再構築と競合する可能性があるため、チェックポイント化された学習中にキーバリューキャッシュを無効にする必要がある場合もあります。

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 peak allocated GPU memoryなどの測定基準を使用して、チェックポイントありとなしの実行を比較する必要があります。

実社会での応用#

• 高解像度コンピュータビジョン: 医療用セグメンテーション、空撮検出、産業用検査では、モデルの重みよりも特徴マップが多くのメモリを消費する大きな画像で学習が行われる場合があります。選択したバックボーンステージをチェックポイント化することで、入力を積極的にダウンスケールする代わりに、画像の解像度を維持したり、バッチあたりのサンプル数を増やしたりすることができます。

• 長シーケンスのTransformer学習: シーケンス長とレイヤー数が増加するにつれて、アクティベーションの保存量急増します。Transformerブロックを再計算することで、同じアクセラレータに長いコンテキストや大きなマイクロバッチを収めることができます。NVIDIA activation recomputation guideは、Transformerレイヤーの完全な再計算と選択的な再計算を示しています。

実践的なガイダンス#

プロファイリングにより、パラメータやオプティマイザの状態ではなく、アクティベーションがメモリを占有していることが判明した場合は、勾配チェックポイントを使用してください。大きな反復ブロックから始め、カバレッジを拡大する前に、ピークメモリ、イテレーション時間、および検証動作をベンチマークします。

文書化されているUltralytics YOLO training configurationsの場合、第一選択のメモリコントロールには、自動混合精度、画像サイズ、物理バッチサイズが含まれます。Ultralytics AutoBatch referenceでは、利用可能なGPUメモリに基づいた自動バッチ選択について説明しています。ローカルハードウェアが不十分な場合、Ultralytics Platform cloud trainingでは、マネージドトレーニング実行用の構成可能なクラウドGPUを提供しています。カスタムPyTorchアーキテクチャでよりきめ細かなアクティベーションメモリ管理が必要な場合、勾配チェックポイントはこれらのコントロールを補完できます。