勾配チェックポインティング（コウバイチェックポインティング）
ニューラルネットワークの訓練時に中間活性値（Activation）を全て保持する代わりに一部のみ保存し、逆伝播時に必要な活性値を再計算することでGPUメモリ使用量を大幅に削減する手法。
勾配チェックポインティングとは
勾配チェックポインティング（Gradient Checkpointing）は、ディープラーニングの訓練時におけるメモリと計算のトレードオフを制御する技術である。2016年にTianqi Chenらが「Training Deep Nets with Sublinear Memory Cost」で提案した手法で、メモリ使用量をO(n)からO(√n)に削減できる（nはネットワークの層数）。
通常のバックプロパゲーションでは、順伝播（Forward Pass）で計算した各層の中間活性値を全てGPUメモリに保持し、逆伝播（Backward Pass）で勾配計算に使用する。100層のネットワークでは100層分の活性値がメモリを占有する。勾配チェックポインティングでは、特定の層（チェックポイント）の活性値のみを保持し、逆伝播時にチェックポイント間の活性値を再計算する。
動作原理
標準的なバックプロパゲーション
順伝播で全層の出力を保持：
- Layer 1 → 活性値a1を保存
- Layer 2 → 活性値a2を保存
- ...
- Layer N → 活性値aNを保存
逆伝播でa1〜aNを参照して勾配を計算。メモリ使用量はO(N)。
チェックポインティング適用後
√N層ごとにチェックポイントを設定（例：100層なら10層ごと）：
- Layer 1 → 活性値a1を保存（チェックポイント）
- Layer 2〜9 → 活性値を破棄
- Layer 10 → 活性値a10を保存（チェックポイント）
- Layer 11〜19 → 活性値を破棄
- ...
逆伝播時、Layer 9の勾配が必要になったらa1から再計算してa2〜a9を得る。メモリ使用量はO(√N)、計算量は約1.33倍に増加。
メモリ削減効果の実測値
| モデル | パラメータ数 | チェックポイントなし | チェックポイントあり | 削減率 |
|---|---|---|---|---|
| GPT-2 Large | 774M | 12.4GB | 4.8GB | 61% |
| Llama 2 7B | 7B | 28GB以上 | 14GB | 50% |
| Llama 2 13B | 13B | 52GB以上 | 24GB | 54% |
| Llama 3 70B | 70B | 280GB以上 | 80GB | 71% |
チェックポイント間隔を狭くするほどメモリ削減効果は大きくなるが、再計算コスト（訓練時間の増加）も大きくなる。実用的な設定ではTransformerブロック単位でチェックポイントを設定し、訓練時間の増加を20〜33%に抑えることが多い。
実装方法
PyTorchでの実装
PyTorchでは torch.utils.checkpoint.checkpoint 関数で特定の層に対してチェックポインティングを適用できる。Transformerモデルの場合、各Transformerブロックをチェックポイント対象にするのが一般的。
Hugging Face Transformersライブラリでは model.gradient_checkpointing_enable() を呼ぶだけで有効化できる。内部的にはTransformerの各レイヤーに対して torch.utils.checkpoint が適用される。
DeepSpeedでの実装
DeepSpeedフレームワークではActivation Checkpointingとして実装されており、ZeRO Stage 1〜3と組み合わせて使用することで、さらに大幅なメモリ削減が可能。パーティション化された活性値チェックポイント（Partitioned Activation Checkpointing）では、チェックポイント自体をGPU間で分割保持する。
JAX/Flaxでの実装
JAXでは jax.checkpoint（旧称 jax.remat）デコレータで関数レベルのチェックポインティングを指定できる。ポリシーベースの設定により、再計算コストの高い操作（Attention計算）は保持し、コストの低い操作（活性化関数）のみ再計算する柔軟な制御が可能。
他のメモリ最適化手法との組み合わせ
| 手法 | メモリ削減対象 | 併用効果 |
|---|---|---|
| 勾配チェックポインティング | 活性値メモリ | 基盤手法 |
| Mixed Precision（FP16/BF16） | 全体（活性値+パラメータ+勾配） | 2倍削減、チェックポイントと独立に適用可 |
| 勾配累積 | バッチメモリ | 実効バッチサイズを増やしつつメモリ抑制 |
| ZeRO-Offload | パラメータ+勾配メモリ | CPUメモリを活用してGPUメモリを解放 |
| FlashAttention | Attention活性値 | Attention部分のメモリをO(N²)→O(N)に削減 |
実運用では複数手法を組み合わせるのが標準。例えばLlama 3 70Bのファインチューニングでは、勾配チェックポインティング + LoRA + BF16 + FlashAttention-2 の組み合わせで、A100 80GB×1枚での学習が可能になる。
FAQ
Q1: 勾配チェックポインティングで訓練精度は低下しますか？
精度は低下しない。活性値を「保持する」か「再計算する」かの違いだけで、計算される勾配の値は数値的に同一である（浮動小数点の丸め誤差を除く）。メモリと計算時間のトレードオフであり、品質のトレードオフではない。
Q2: チェックポイントの間隔はどう決めるべきですか？
理論的にはO(√N)間隔が最適だが、実用上はTransformerブロック単位（1ブロックごと）が最もよく使われる。GPUメモリに余裕がある場合は2〜4ブロックごとに設定して再計算コストを抑え、メモリが逼迫している場合は全ブロックに適用する。
Q3: 推論時にも勾配チェックポインティングは使えますか？
推論時は逆伝播が不要なため、勾配チェックポインティングは適用されない。推論時のメモリ最適化にはKVキャッシュの量子化、PagedAttention、動的バッチングなど別の手法が用いられる。