啟用卸載

一般而言，向前傳遞會計算每一層的啟用，並將其保留在 GPU 記憶體中，直到對應層的向後傳遞完成為止。在向前傳遞之後將這些張量卸載到 CPU 記憶體，並在需要時將張量擷取回 GPU，可以節省大量的 GPU 記憶體使用量。PyTorch 支援卸載啟用，但實作會導致 GPU 在回溯傳遞期間從 CPU 擷取啟用時閒置。使用啟用卸載時，這會導致主要效能降低。

SMP v2 改善了此啟用卸載。它會在需要啟用之前預先擷取啟用，GPU 才能開始向後傳遞這些啟用。預先擷取功能有助於更有效率地執行訓練進度，無需閒置 GPU。這會導致記憶體使用率降低，而不會降低效能。

您可以保留原生 PyTorch 模組，以便在訓練指令碼中卸載啟用。以下是在指令碼中套用 SMP 啟用卸載功能的範例結構。請注意，啟用卸載「僅」適用於與 啟用檢查點 搭配使用時。若要進一步了解用於啟用卸載的原生 PyTorch 檢查點工具，請參閱：

您可以在 PyTorch 啟用檢查點上套用 SMP 啟用卸載功能。這可透過在步驟 2：啟動訓練任務 期間將 sm_activation_offloading 和 activation_loading_horizon 參數新增至 SMP 組態字典來完成。

下列程式碼片段示範如何在訓練指令碼中新增 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)