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

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

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

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

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

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

SMP 設定

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

トレーニングスクリプト

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)