

本文為英文版的機器翻譯版本，如內容有任何歧義或不一致之處，概以英文版為準。

# 啟用卸載
<a name="model-parallel-core-features-v2-pytorch-activation-offloading"></a>

**重要**  
在 SMP v2.2.0 中，SMP 程式庫的啟用卸載功能無法運作。請改用原生 PyTorch 啟用卸載。

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

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

您可以保留原生 PyTorch 模組，以便在訓練指令碼中卸載啟用。以下是在指令碼中套用 SMP 啟用卸載功能的範例結構。請注意，啟用卸載**「僅」適用於與 [啟用檢查點](model-parallel-core-features-v2-pytorch-activation-checkpointing.md) 搭配使用時。若要進一步了解用於啟用卸載的原生 PyTorch 檢查點工具，請參閱：
+ *PyTorch GitHub 儲存庫*中的 [checkpoint\_wrapper.py](https://github.com/pytorch/pytorch/blob/v2.0.1/torch/distributed/algorithms/_checkpoint/checkpoint_wrapper.py#L171)
+ PyTorch 部落格「使用 PyTorch 分散式擴展 TorchMultimodal 中的多模態基礎模型」**中的[啟用檢查點](https://pytorch.org/blog/scaling-multimodal-foundation-models-in-torchmultimodal-with-pytorch-distributed/#activation-checkpointing)。

您可以在 [PyTorch 啟用檢查點](https://pytorch.org/blog/scaling-multimodal-foundation-models-in-torchmultimodal-with-pytorch-distributed/#activation-checkpointing)上套用 SMP 啟用卸載功能。這可透過在[步驟 2：啟動訓練任務](model-parallel-use-api-v2.md#model-parallel-launch-a-training-job-v2) 期間將 `sm_activation_offloading` 和 `activation_loading_horizon` 參數新增至 SMP 組態字典來完成。

下列程式碼片段示範如何在訓練指令碼中新增 SMP 初始化模組 `torch.sagemaker.init()`，以及設定訓練任務啟動器的 JSON 格式 SMP 組態字典，同時遵循[使用 SageMaker 模型平行化程式庫 v2](model-parallel-use-api-v2.md) 中介紹的兩個步驟程序。您不需要對 PyTorch 模型或 [PyTorch FSDP](https://pytorch.org/docs/stable/fsdp.html#module-torch.distributed.fsdp) 組態進行任何變更。如需 `sm_activation_offloading` 和 `activation_loading_horizon` 參數的詳細資訊，請參閱 [SMP v2 核心功能組態參數](distributed-model-parallel-v2-reference.md#distributed-model-parallel-v2-reference-init-config)。

**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)
```