Activation Checkpointing（アクティベーション チェックポインティング）
深層学習の学習時に中間活性化値をすべてメモリに保持する代わりに、一部のレイヤーの出力のみを保存し、逆伝播時に必要な活性化値を再計算する手法。GPUメモリ使用量を大幅に削減し、より大きなモデルやバッチサイズでの学習を可能にする。
Activation Checkpointing とは
Activation Checkpointing（勾配チェックポインティング、Gradient Checkpointing とも呼ばれる）は、ニューラルネットワークの学習時における GPU メモリ消費を削減するための手法です。通常、順伝播（Forward Pass）で計算した各レイヤーの中間活性化値はすべてメモリに保持され、逆伝播（Backward Pass）で勾配計算に使用されます。大規模モデルではこの活性化値のメモリ消費がモデルパラメータ自体の数倍に達するため、学習のボトルネックとなります。
メモリ問題の背景
Transformer ベースの LLM では、活性化値のメモリ消費は以下の要因で決まります。
| パラメータ | メモリへの影響 |
|---|---|
| バッチサイズ (B) | 線形に増加 |
| シーケンス長 (S) | 線形〜二乗に増加（Self-Attention は S² に比例） |
| 隠れ層の次元 (H) | 線形に増加 |
| レイヤー数 (L) | 線形に増加 |
例えば、Llama 2 7B（32層、H=4096）をシーケンス長 4096、バッチサイズ 4 で学習する場合、活性化値だけで約 60GB のメモリを消費します。モデルパラメータ（FP16 で約 14GB）とオプティマイザ状態（Adam で約 56GB）を加えると、合計 130GB 以上が必要です。
動作原理
Activation Checkpointing は「メモリと計算のトレードオフ」を実現します。
チェックポイントなし（標準）: 全 L 層の活性化値を保持。メモリ消費 O(L)、再計算なし。
全レイヤーチェックポイント: 特定の間隔（例えば √L 層ごと）にチェックポイントを設定。チェックポイント間の活性化値は破棄し、逆伝播時に再計算。メモリ消費 O(√L)、順伝播の約 33% を再計算。
セグメント単位チェックポイント: Transformer ブロック単位でチェックポイントを設定。各ブロックの入力のみを保存し、ブロック内部の活性化値は逆伝播時に再計算。
実装方式の比較
| 方式 | メモリ削減率 | 計算オーバーヘッド | 実装の容易さ |
|---|---|---|---|
| チェックポイントなし | 0% | 0% | - |
| 全レイヤー均等 | 60〜70% | 30〜35% | 容易 |
| セグメント（ブロック）単位 | 50〜60% | 25〜30% | 中程度 |
| 選択的チェックポイント | 40〜80% | 15〜25% | 要分析 |
選択的チェックポイント: メモリ消費量が大きいレイヤー（Self-Attention の QKV 行列など）のみチェックポイントし、メモリ消費が小さいレイヤー（LayerNorm など）は保持する方式。分析コストがかかるが、最適なメモリ/計算トレードオフを実現します。
PyTorch での実装
PyTorch では torch.utils.checkpoint.checkpoint 関数で簡単に利用できます。
Hugging Face Transformers では model.gradient_checkpointing_enable() を呼ぶだけで有効化されます。DeepSpeed や FSDP などの分散学習フレームワークでも設定一つで有効化可能です。
他のメモリ最適化手法との組み合わせ
Activation Checkpointing は単体でも効果的ですが、他の手法と組み合わせることでさらなるメモリ削減が可能です。
- 混合精度学習（Mixed Precision）: FP16/BF16 で活性化値のサイズを半減。チェックポイントと併用で 70〜80% 削減
- ZeRO（Zero Redundancy Optimizer）: オプティマイザ状態のメモリを分散。チェックポイントは活性化値を削減
- メモリオフロード: GPU メモリから CPU メモリへ一部をオフロード。チェックポイントで残りを削減
- Gradient Accumulation: バッチサイズを論理的に拡大しつつ物理メモリ消費を抑制
これらを全て組み合わせることで、RTX 4090（24GB VRAM）でも 13B パラメータモデルのフルファインチューニングが可能になります。
FAQ
Q1: Activation Checkpointing は推論時にも使えますか？
A1: 推論時は逆伝播が不要なため、活性化値を保持する必要がなく、チェックポインティングは無意味です。推論のメモリ最適化には KV Cache の量子化やページングなど別の手法を使います。
Q2: 計算オーバーヘッド 30% は学習時間に直結しますか？
A2: 理論上は順伝播の再計算分だけ遅くなりますが、実際にはメモリ削減によりバッチサイズを大きくできるため、スループット（トークン/秒）はほぼ維持できるケースが多いです。A100 80GB でチェックポイントなし・バッチ4 vs チェックポイントあり・バッチ8 なら後者の方が高速です。
Q3: LoRA などの PEFT 手法でもチェックポイントは必要ですか？
A3: LoRA は学習パラメータを大幅に削減しますが、順伝播の活性化値は全モデルで計算されるためメモリ消費は残ります。7B モデルの LoRA ファインチューニングでもチェックポイントを有効化すると VRAM 使用量が 4〜6GB 削減されます。