激活分载 - Amazon SageMaker
如何使用激活分载
Amazon Web Services 文档中描述的 Amazon Web Services 服务或功能可能因区域而异。要查看适用于中国区域的差异，请参阅 中国的 Amazon Web Services 服务入门 (PDF)。

激活分载

当激活检查点和管道并行性均已启用并且微批次数量大于 1 时，激活分载是可以进一步减少内存使用量的附加功能。对于当前未在 CPU 中运行的微批次，激活分载会异步移动与其对应的有检查点的激活。就在 GPU 需要激活以便微批次向后传递之前，此功能会将分载的激活从 CPU 预取回来。

注意

此功能在 SageMaker 模型并行库 v1.6.0 及更高版本上对 PyTorch 可用。

如何使用激活分载

当微批次的数量大于 1 并启用了激活检查点时，请使用激活分载以减少内存使用量（请参阅激活检查点）。如果不使用激活检查点，则激活分享不会生效。在仅与一个微批次一起使用时，这不会节省内存。

要使用激活分载，请在 modelparallel 配置中设置 "offload_activations": True。

激活分载将 nn.Sequential 模块中的有检查点的激活异步移至 CPU。通过 PCIe 链路传输的数据与 GPU 计算重叠。在计算了特定检查点层的向前传递后，会立即进行分载。在特定微批次的向后传递需要激活之前的片刻，激活将加载回 GPU。CPU-GPU 传输同样与计算重叠。

要调整激活加载回 GPU 的时间提前量，您可以使用配置参数 "activation_loading_horizon"（默认设置为 4，必须为大于 0 的 int）。较大的激活加载范围会导致激活更早地加载回 GPU。如果范围太大，激活分载所带来的内存节省影响可能会减弱。如果范围太小，激活可能无法及时加载回，从而减少重叠量，导致性能下降。

提示

激活分载对于参数超过一千亿的大型模型非常有用。

配置 SageMaker PyTorch 估算器

mpi_options = {
    "enabled" : True,
    "processes_per_host" : 8,               # 8 processes
    "custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none "
}

smp_options = {
    "enabled":True,
    "parameters": {
        "microbatches": 4,
        "pipeline_parallel_degree": 2,    # alias for "partitions"
        "placement_strategy": "cluster",
        "tensor_parallel_degree": 2,      # tp over 2 devices
        "ddp": True,
        "offload_activations": True,
        "activation_loading_horizon": 4   # optional. default is 4.
    }
}