勾配チェックポイントとは
勾配チェックポイント（Gradient Checkpointing、Activation Checkpointing）は、ディープラーニングの学習時にGPUメモリの使用量を削減するための技術です。通常、逆伝播（バックプロパゲーション）では順伝播時の中間活性化値をすべてメモリに保持する必要がありますが、勾配チェックポイントでは一部の活性化値のみを保持し、必要に応じて再計算することでメモリを節約します。
| ひとことで言うと | 途中の計算結果を保存せず必要時に再計算し、メモリを大きく節約する技術。 |
|---|---|
| どこで使う | GPUメモリが足りず、大きなモデルや長い系列を学習したいとき。 |
| 注意点 | メモリは減るが計算時間は2〜3割増える。時間とメモリのトレードオフ。 |
メモリと計算のトレードオフ
勾配チェックポイントの本質は、メモリと計算時間のトレードオフです。チェックポイントを設定した層の活性化値のみを保存し、その間の層の活性化値は逆伝播時に再計算します。これにより、メモリ使用量はO(n)からO(√n)に削減できますが、順伝播の計算が約1.3〜1.5倍増加します。大規模モデルではメモリがボトルネックになることが多いため、この追加計算は許容されることが一般的です。
1通常は全部保存逆伝播に備えて記録
▶
2一部だけ保存要所（チェックポイント）のみ
▶
3逆伝播で再計算捨てた部分を作り直す
▶
4メモリが大幅減大きなモデルが載る
実装方法
PyTorchではtorch.utils.checkpoint.checkpoint関数を使用して、特定のモジュールの活性化値の保存をスキップできます。TensorFlowではtf.recompute_gradが同等の機能を提供します。モデル全体に適用することも、メモリ消費が大きい特定の層のみに適用することもできます。Transformerの各層にチェックポイントを設定するのが一般的なアプローチです。
大規模モデル学習での重要性
GPTやLLaMAなどの大規模言語モデルの学習では、勾配チェックポイントは不可欠な技術です。数百億パラメータのモデルでは、すべての活性化値を保持するとGPUメモリが不足するため、勾配チェックポイントに加えて、モデル並列化やZeRO最適化と組み合わせてメモリ効率を最大化します。