アクティベーションオフロード

重要

SMP v2.2.0 では、SMP ライブラリのアクティベーションオフロード機能は機能しません。代わりにネイティブ PyTorch アクティベーションオフロードを使用します。

通常、フォワードパスは各レイヤーのアクティベーションを計算し、対応するレイヤーのバックワードパスが終了するまで GPU メモリに保持します。これらのテンソルをフォワードパス後に CPU メモリにオフロードし、必要なときに GPU にフェッチバックすると、GPU メモリの使用量を大幅に節約できます。 はアクティベーションのオフロード PyTorch をサポートしますが、この実装により、アクティベーションがバックワードパス中に CPU からフェッチバックされる間に GPUs がアイドル状態になります。これにより、アクティベーションオフロードの使用時にパフォーマンスが大幅に低下します。

SMP v2 は、このアクティベーションのオフロードを改善します。GPU がアクティベーションのバックワードパスを開始するには、アクティベーションが必要になる前にアクティベーションをプリフェッチします。プリフェッチ機能は、アイドル状態の GPUs を使用せずにトレーニングの進行状況をより効率的に実行するのに役立ちます。これにより、パフォーマンスを低下させることなく、メモリ使用量を減らすことができます。

ネイティブ PyTorch モジュールは、トレーニングスクリプトでアクティベーションをオフロードするために保持できます。以下は、スクリプトで SMP アクティベーションオフロード機能を適用する構造の例です。アクティベーションオフロードは、 と一緒に使用する場合にのみ適用されることに注意してくださいアクティベーションチェックポイント。アクティベーションオフロード用のネイティブ PyTorch チェックポイントツールの詳細については、以下を参照してください。

• PyTorch GitHub リポジトリ内の checkpoint_wrapper.py

• PyTorch ブログ「分散 を使用した でのマルチモーダル基盤モデルのスケーリング」の「アクティベーションチェックポイント」。 TorchMultimodal PyTorch

SMP アクティベーションオフロード機能は、PyTorch アクティベーションチェックポイント に適用できます。これは、 中に sm_activation_offloadingおよび activation_loading_horizonパラメータを SMP 設定ディクショナリに追加することによって行われますステップ 2: トレーニングジョブを起動する。

次のコードスニペットは、トレーニングスクリプトtorch.sagemaker.init()に SMP 初期化モジュールを追加し、 で導入された 2 ステッププロセスに従って、トレーニングジョブランチャーの JSON 形式で SMP 設定ディクショナリを設定する方法を示しています SageMaker モデル並列処理ライブラリ v2 を使い始める。 PyTorch モデルまたは PyTorch FSDP 設定を変更する必要はありません。sm_activation_offloading、および activation_loading_horizon、パラメータの詳細については、SMP v2 コア機能の構成パラメーター を参照してください。

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)