Gradient checkpointing

概要

学習ステップは2つの段階に分かれます。フォワードパスでは、各レイヤーが入力から出力を生成し、その過程で一連の中間結果であるアクティベーションを生成します。次に、バックワードパスがレイヤーを逆方向に処理して勾配を計算しますが、各レイヤーの勾配にはそのレイヤーのアクティベーションが必要になります。単純なアプローチでは、バックワードパスが到達するまでそれらをすべてメモリに保持します。長いサンプルの場合、アクティベーションはモデル自体の重みよりも多くのメモリを消費する可能性があります。

Gradient checkpointing（アクティベーション・チェックポイントまたは再計算とも呼ばれます）は、一部のアクティベーション（チェックポイント）のみを保持し、残りを破棄する手法です。バックワードパスで破棄されたアクティベーションが必要になった場合、最も近いチェックポイントからフォワードパスを再実行して再作成します。Chenら（2016年）は、バッチごとに約1回の追加フォワードパスのコストで、n層ネットワークのアクティベーションメモリを√nのオーダーまで削減できることを示しました。彼らの実験では、1,000層のネットワークで実行時間が約30%増加したものの、メモリ使用量は48GBから7GBに減少しました。

言語モデルにおける仕組み

Transformerライブラリでは、通常、チェックポイントは層の境界に配置されます。各層の入力は保持されますが、層の内部で計算されたものはすべて破棄され、バックワードパスがそこに到達した際に再構築されます。そのため、すべてのアクティベーションが常に存在するのではなく、ある瞬間には1つの層の完全なアクティベーションセットのみが存在することになります。

activations per GPT-style layer, 16-bit, efficient attention   (Korthikanti et al., 2022)
  ≈ 34 × batch × tokens × hidden size   bytes

illustrative 2B-class model: hidden size 2,048 · 28 layers · batch 2 × 1,024 tokens
  all kept                 34 × 2 × 1,024 × 2,048 × 28   ≈ 4.0 GB
  checkpointed, inputs      2 × 2 × 1,024 × 2,048 × 28   ≈ 0.23 GB
  + one layer rebuilt at a time                          ≈ 0.14 GB

保持するチェックポイントの数は調整可能です。多く保持すれば再計算は減りますがメモリを多く消費し、少なく保持すればその逆になります。すべての層の境界でチェックポイントを設定するのが通常のデフォルトです。これはシンプルでありながら、アクティベーションメモリの大部分をすでに削減できるためです。残るのは、各層の小さな入力1つと、再構築中の1つの層のみです。

その代償となるのは計算量です。フォワードパスは実質的に2回実行されます。1回目は通常通り、2回目はバックワードパスの間に少しずつ実行されます。バックワードパスはすでにフォワードパスの約2倍のコストがかかるため、追加の作業は最大でも1ステップの約3分の1程度です。実際のところ、トレーニングの実行速度は約20〜30%遅くなります。結果は変わりません。再計算された値は同じものであるため、微小な数値の違いを除けば、損失曲線はチェックポイントを使用しない実行と一致します。

大規模な語彙において重要性が増す理由

モデルの最後では、すべてのサンプルのすべての位置において、語彙内のすべてのトークンに対するスコア、すなわちロジットが算出されます。そのサイズはモデルの深さには依存せず、トレーニングされる位置の数と語彙の大きさにのみ依存します。新しいモデルでは150,000から250,000トークンの語彙を持つため、ロジットだけで数ギガバイトを消費する可能性があります。

logits = batch × tokens × vocabulary size × bytes per number

2 × 1,024 × 248,000 × 4 bytes (32-bit)   ≈ 2.0 GB   before their gradients

効果的な活用方法

• データを削減する前に検討してください。 チェックポイントを有効にすると時間がかかりますが、最大長を短くすると長いサンプルの末尾が切り捨てられます。メモリが限界に達している場合、通常は速度低下を受け入れる方が良い選択です。
• 解放されたメモリを活用してください。 バッチサイズを大きくしたり、最大長を長くしたりすることで、エポックごとに失われた速度の一部を取り戻すことができます。
• マシンが保持できる最大規模のモデルには、QLoRAと組み合わせてください。QLoRAは重みを縮小し、チェックポイントはアクティベーションを縮小します。
• 影響を与えないものを理解してください。 重み、勾配、オプティマイザの状態は同じサイズのままであり、完成したモデルの実行には何の影響も与えません。
• デフォルトを確認してください。 多くのファインチューニングのレシピでは、すでに有効になっています。メモリに余裕があり、速度が重要な場合は、無効にすることで各ステップを高速化できます。

よくある間違い

• チェックポイントのみでメモリに収まるにもかかわらず、最初に最大長を短くし、長いサンプルの末尾を切り捨ててしまうこと。
• 重みやオプティマイザの状態が縮小されると期待すること。影響を受けるのはアクティベーションのみです。
• 20〜30%の速度低下をマシンの不具合と解釈すること。これはメモリを節約するための代償です。
• ロジットのコストが排除されると思い込むこと。膨大な語彙がある場合、依然としてより小さなバッチや最大長が必要になることがあります。

参考文献および関連資料

外部ソースは新しいタブで開きます。

1. 01Training Deep Nets with Sublinear Memory Cost↗Chen et al., 2016 · arXiv · arxiv.org再計算を標準的な手法にした論文：1回の追加フォワードパスでメモリを√nに削減。
2. 02Reducing Activation Recomputation in Large Transformer Models↗Korthikanti et al., 2022 · arXiv · arxiv.org層ごとのアクティベーション見積もりの由来と、コストの高い部分のみを再計算する方法。
3. 03torch.utils.checkpoint↗PyTorch docs · docs.pytorch.org
4. 04GPU memory usage↗Hugging Face Transformers docs · huggingface.co重み、勾配、オプティマイザの状態、アクティベーション：トレーニングメモリの消費先。

関連する概念