啟用卸載

重要

在 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：啟動訓練任務。

下列程式碼片段說明如何將SMP初始化模組新增至torch.sagemaker.init()訓練指令碼，並針對訓練任務啟動器設定JSON格式的SMP組態字典，同時遵循 中介紹的兩步驟程序使用 SageMaker 模型平行處理程式庫 v2。您不需要對 PyTorch 模型或PyTorch FSDP組態進行任何變更。如需 sm_activation_offloading 和 activation_loading_horizon 參數的詳細資訊，請參閱 SMPv2 核心功能配置參數。

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)