激活分载

通常情况下，前向传递计算每一层的激活度，并将其保存在 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()，并按照 使用 SageMaker 模型并行度库 v2 中介绍的两步流程，为训练作业启动器设置 JSON 格式的 SMP 配置字典。您无需对 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)