# 勾配チェックポイント

フォワードパスでは、バックワードパスで再利用するためにすべての中間アクティベーションがキャッシュされるのが一般的です。しかし、アクティベーションはバッチサイズとシーケンス長に比例して増加します。Gradient checkpointingでは、一部のアクティベーションのみを保存し、残りは破棄します。これにより、バックワードパスでは必要に応じて一部のアクティベーションをオンザフライで再計算する必要があります。

```text
Normal training:
  Forward:   [L1]→[L2]→[L3]→[L4]   (save ALL activations)
  Backward:  ←uses cached activations everywhere

Gradient checkpointing:
  Forward:   [L1]→[L2]→[L3]→[L4]   (save only at checkpoints, discard the rest)
  Backward:  ←reaches L2, recomputes L2→L3 from scratch, uses it, discards it
```

バックワードパスで破棄されたアクティベーションが再計算されるため、通常はトレーニングが遅くなりますが、チェックポイント化によりアクティベーションのメモリ使用量を削減できます。

`gradient_checkpointing=True` を設定して有効にします。

> [!TIP]
> [勾配蓄積](./grad_accumulation)と組み合わせて使用し、メモリ使用量をさらに削減できます。

```py
from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    gradient_checkpointing=True,
)
```

## 部分的なチェックポイント

フルのグラデーション・チェックポイントでは、チェックポイント可能なすべてのレイヤーを再計算します。実行にメモリに余裕がある場合は、チェックポイントするレイヤーを減らし、メモリ削減の一部を速度向上に振り分けることができます。

`every_n_layers` を [gradient_checkpointing_enable()](/docs/transformers/main/ja/main_classes/model#transformers.PreTrainedModel.gradient_checkpointing_enable) に渡すことで、チェックポイントの間隔を選択できます。

```text
every_n_layers=2

Forward:   input  -> [L1] -> [L2] -> [L3] -> [L4] -> [L5] -> [L6]
                      CP     keep     CP     keep     CP     keep

Backward:  output <- [L6] <- [L5] <- [L4] <- [L3] <- [L2] <- [L1]
                     keep   rerun    keep   rerun    keep   rerun
```

`every_n_layers=2`を指定すると、最初のレイヤーと、その後の2つおきのレイヤーがチェックポイント化されます。チェックポイント化されたレイヤーは、フォワードパス中にアクティベーションを破棄し、バックワードパス中に再計算します。一方、その他のレイヤーはアクティベーションをメモリに保持します。

```py
model.gradient_checkpointing_enable(every_n_layers=2)
```

デフォルトの`every_n_layers=1`では、すべてのレイヤーでチェックポイントが作成されます。より大きな値を指定すると、最初のレイヤーでチェックポイントを作成し、その後、`n`レイヤーごとにチェックポイントを作成します。その他のレイヤーのアクティベーションはメモリ内に保持されます。例えば、`every_n_layers=2`を指定すると、レイヤー 1、3、5 などがチェックポイントされます。カウントされるのは、[GradientCheckpointingLayer](/docs/transformers/main/ja/internal/modeling_utils#transformers.GradientCheckpointingLayer)を継承するモジュールのみです。勾配チェックポイントに対応するその他のモジュールは、引き続き有効なままです。

[Trainer](/docs/transformers/main/ja/main_classes/trainer#transformers.Trainer)で部分的な勾配チェックポイントを使用するには、`gradient_checkpointing_kwargs`内の`every_n_layers`を設定します。

```py
from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"every_n_layers": 4},
)
```

  

## 保存されたアクティベーションのオフロード

Gradient checkpointing は、チェックポイント化された各レイヤーごとに1つのアクティベーションを GPU に保持します。シーケンス長が長い場合、これにより大量のメモリを消費する可能性があります（約 `layers x sequence x hidden x bytes_per_element`）。`offload` を設定すると、これらのアクティベーションをピン留めされたホストメモリに保持できます。その代わりとして、フォワードパス時にデバイスからホストへのコピー、バックワードパス時にホストからデバイスへのコピーが発生します。

```py
from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"offload": True},
)
```

両方のコピーは compute stream 上で実行されるため、メモリを消費する代わりにステップが遅くなります。実行がそれ以外の方法では収まらない場合に使用してください。高速化を目的として使用しないでください。

## 次のステップ

- GPU上のメモリ使用量を駆動している要因を理解するために、[GPUメモリ使用量](./model_memory_anatomy)のドキュメントをお読みください。
- メモリをさらに削減し、トレーニングを高速化するために、低精度のデータ型を使用する方法については、[Mixed precision training](./mixed_precision_training)ガイドを参照してください。
- カスタム融合カーネルによる学習の高速化方法については、[Kernels](./kernels)ガイドを参照してください。