

本文属于机器翻译版本。若本译文内容与英语原文存在差异，则一律以英文原文为准。

# 激活分载


**重要**  
在 SMP v2.2.0 中，SMP 库的激活卸载功能不起作用。改用本机 PyTorch 激活卸载。

通常情况下，前向传递计算每一层的激活度，并将其保存在 GPU 内存中，直到相应层的后向传递结束。在向前传递后将这些张量卸载到 CPU 内存中，并在需要时将其取回 GPU，可以节省大量的 GPU 内存使用量。 PyTorch 支持卸载激活，但是实现会导致在 GPUs 向后传递期间从 CPU 获取激活时处于空闲状态。在使用激活卸载时，这会导致性能严重下降。

SMP v2 改进了这种激活卸载。在 GPU 开始后向传递这些激活信息之前，它会提前获取激活信息。预取功能有助于在不闲置的情况下更高效地运行训练进度。 GPUs这样既可以降低内存使用量，又不会降低性能。

您可以在训练脚本中保留用于卸载激活的本机 PyTorch 模块。以下是在脚本中应用 SMP 激活卸载功能的结构示例。请注意，激活卸载*仅*在与 [激活检查点](model-parallel-core-features-v2-pytorch-activation-checkpointing.md) 一起使用时才适用。要了解有关用于激活卸载的本机 PyTorch 检查点工具的更多信息，请参阅：
+ *PyTorch GitHub仓库*里@@ [的 checkpoint\$1wrapper.py](https://github.com/pytorch/pytorch/blob/v2.0.1/torch/distributed/algorithms/_checkpoint/checkpoint_wrapper.py#L171)
+  PyTorch 博客 “使用分布式*扩展多模态基础模型” 中的 TorchMultimodal *[激活检查点](https://pytorch.org/blog/scaling-multimodal-foundation-models-in-torchmultimodal-with-pytorch-distributed/#activation-checkpointing)。 PyTorch 

[您可以在激活检查点上PyTorch 应用 SMP 激活卸载功能。](https://pytorch.org/blog/scaling-multimodal-foundation-models-in-torchmultimodal-with-pytorch-distributed/#activation-checkpointing)具体做法是在 [步骤 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()`，并按照 [使用 SageMaker 模型并行度库 v2](model-parallel-use-api-v2.md) 中介绍的两步流程，为训练作业启动器设置 JSON 格式的 SMP 配置字典。您无需对 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)
```