激活分载

重要

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

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

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

您可以在训练脚本中保留用于卸载激活的本机 PyTorch 模块。以下是在脚本中应用 SMP 激活卸载功能的示例结构。请注意，激活卸载仅在与激活检查点一起使用时才适用。要了解有关用于激活卸载的本机 PyTorch 检查点工具的更多信息，请参阅：

• PyTorch GitHub仓库里@@ 的 checkpoint_wrapper.py

• PyTorch 博客 “使用分布式扩展多模态基础模型” 中的 TorchMultimodal 激活检查点。 PyTorch

您可以在激活检查点上PyTorch 应用 SMP 激活卸载功能。这是通过在期间步骤 2：启动训练作业向 SMP 配置字典中添加sm_activation_offloading和activation_loading_horizon参数来完成的。

以下代码片段展示了如何将 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
}

在训练脚本中

注意

在激活 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)