

翻訳は機械翻訳により提供されています。提供された翻訳内容と英語版の間で齟齬、不一致または矛盾がある場合、英語版が優先します。

# アクティベーションオフロード
<a name="model-parallel-core-features-v2-pytorch-activation-offloading"></a>

**重要**  
SMP v2.2.0 では、SMP ライブラリのアクティベーションオフロード機能が動作しません。代わりに、ネイティブの PyTorch アクティベーションオフロードを使用してください。

通常、フォワードパス (順伝播) では、各層でアクティベーションが計算され、対応する層のバックワードパス (逆伝播) が終了するまで、それらが GPU メモリに保持されます。これらのテンソルをフォワードパスの後に CPU メモリにオフロードし、必要になった時点で GPU に戻すことで、GPU メモリの使用量を大幅に削減できます。PyTorch はアクティベーションのオフロードに対応していますが、その実装では、バックワードパス中に CPU から GPU にアクティベーションを戻す際に、GPU がアイドル状態になります。そのせいで、アクティベーションオフロードを使用すると、パフォーマンスが大幅に低下します。

SMP v2 は、このアクティベーションオフロードを改善します。GPU がアクティベーションについてバックワードパスを開始する前に、必要なアクティベーションを前もって取得しておきます。このプリフェッチ機能のおかげで、トレーニングプロセスが効率化し、GPU がアイドル状態にならなくなります。その結果、メモリ使用量を削減しつつ、パフォーマンスの低下を回避できます。

トレーニングスクリプトで、ネイティブの PyTorch モジュールをそのままアクティベーションオフロードに使用できます。以下は、スクリプトで SMP のアクティベーションオフロード機能を適用する場合の構造例です。アクティベーションオフロードは、[アクティベーションチェックポイント](model-parallel-core-features-v2-pytorch-activation-checkpointing.md) と組み合わせて使用する場合に*のみ*適用可能です。アクティベーションオフロード用のネイティブ PyTorch チェックポイントツールの詳細については、以下を参照してください。
+ *PyTorch GitHub リポジトリ*の [checkpoint\$1wrapper.py](https://github.com/pytorch/pytorch/blob/v2.0.1/torch/distributed/algorithms/_checkpoint/checkpoint_wrapper.py#L171)
+ PyTorch ブログ「*Scaling Multi-modal Foundation Models in TorchMultimodal with PyTorch Distributed*」の「[Activation Checkpointing](https://pytorch.org/blog/scaling-multimodal-foundation-models-in-torchmultimodal-with-pytorch-distributed/#activation-checkpointing)」

SMP のアクティベーションオフロード機能を [PyTorch のアクティベーションチェックポイント](https://pytorch.org/blog/scaling-multimodal-foundation-models-in-torchmultimodal-with-pytorch-distributed/#activation-checkpointing)に適用できます。その場合は、「[ステップ 2: トレーニングジョブを開始する](model-parallel-use-api-v2.md#model-parallel-launch-a-training-job-v2)」の間に、SMP 設定ディクショナリに `sm_activation_offloading` パラメータと `activation_loading_horizon` パラメータを追加します。

次のコードスニペットでは、「[SageMaker モデル並列処理ライブラリ v2 を使用する](model-parallel-use-api-v2.md)」で紹介した 2 ステップのプロセスに従って、SMP 初期化モジュール `torch.sagemaker.init()` をトレーニングスクリプトに追加し、トレーニングジョブランチャーの SMP 設定ディクショナリを JSON 形式で設定する方法を示しています。PyTorch モデルや [PyTorch FSDP](https://pytorch.org/docs/stable/fsdp.html#module-torch.distributed.fsdp) 設定については、一切変更する必要はありません。`sm_activation_offloading`、および `activation_loading_horizon`、パラメータの詳細については、[SMP v2 の主要機能の設定パラメータ](distributed-model-parallel-v2-reference.md#distributed-model-parallel-v2-reference-init-config) を参照してください。

**SMP の設定**

```
{
    "activation_loading_horizon": 2,
    "sm_activation_offloading": True
}
```

**トレーニングスクリプト内**

**注記**  
SMP のアクティベーションオフロード機能を有効にするときは、PyTorch の `offload_wrapper` 関数も併用し、ルートモジュールに適用してください。SMP のアクティベーションオフロード機能は、ルートモジュールを使用して、フォワードパスがいつ完了するかを判断し、完了した時点でプリフェッチを開始します。

```
import torch.sagemaker as tsm
tsm.init()

# Native PyTorch module for activation offloading
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    apply_activation_checkpointing, 
    offload_wrapper,
)

model = FSDP(...)

# Activation offloading requires activation checkpointing.
apply_activation_checkpointing(
    model,
    check_fn=checkpoint_transformer_layers_policy,
)

model = offload_wrapper(model)
```