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

アクティベーションチェックポイントとパイプライン並列処理がオンで、マイクロバッチの数が 1 より大きい場合、アクティベーションオフロードは、メモリ使用量をさらに削減できる追加機能です。アクティベーションオフロードは、CPU で現在実行されていないマイクロバッチに対応するチェックポイントされたアクティベーションを非同期的に移動します。GPU がマイクロバッチのバックワードパスのアクティベーションを必要とする直前に、この機能はオフロードされたアクティベーションを CPU からプリフェッチして戻します。

アクティベーションオフロードの使用方法

マイクロバッチの数が 1 を超え、アクティベーションチェックポイントがオンになっている場合は、アクティベーションオフロードを使用してメモリ使用量を減らします (アクティベーションチェックポイント を参照)。アクティベーションチェックポイントを使用しない場合、アクティベーションオフロードは効果がありません。マイクロバッチを 1 つだけ使用する場合、メモリの節約にはなりません。

アクティベーションオフロードを使用するには、modelparallel 設定で "offload_activations": True を設定してください。

アクティベーションオフロードは、nn.Sequential モジュール内のチェックポイントしたアクティベーションを非同期的に CPU に移動します。PCIe リンクを介したデータ転送は GPU 計算と重複します。オフロードは、特定のチェックポイントしたレイヤーのフォワードパスが計算されるとすぐに行われます。アクティベーションは、特定のマイクロバッチのバックワードパスに必要になる直前に GPU に読み込まれます。CPU - GPU 間の転送も同様に計算と重複します。

アクティベーションが GPU に読み込まれるまでの時間を調整するには、設定パラメータ "activation_loading_horizon" を使用できます (デフォルトは 4 に設定されていますが、0 より int 大きい必要があります)。アクティベーションのロード範囲が大きいと、アクティベーションは GPU に早期にロードバックされます。範囲が大きすぎると、アクティベーションオフロードによるメモリ節約効果が低下する可能性があります。範囲が小さすぎると、アクティベーションが時間内にロードされず、重複の量が減少し、パフォーマンスが低下することがあります。

SageMaker PyTorch 推定器を設定する

mpi_options = {
    "enabled" : True,
    "processes_per_host" : 8,               # 8 processes
    "custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none "
}

smp_options = {
    "enabled":True,
    "parameters": {
        "microbatches": 4,
        "pipeline_parallel_degree": 2,    # alias for "partitions"
        "placement_strategy": "cluster",
        "tensor_parallel_degree": 2,      # tp over 2 devices
        "ddp": True,
        "offload_activations": True,
        "activation_loading_horizon": 4   # optional. default is 4.
    }
}