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

优化器状态分片

优化器状态分片是一种非常有用的内存节省技术，它可以跨数据并行设备组对优化器状态（描述优化器状态的一组权重）进行分片。只要使用有状态优化器（例如 Adam）或 FP16 优化器（它同时存储参数的 FP16 和 FP32 副本），都可以使用优化器状态分片。

注意

优化器状态分片在 SageMaker 模型并行库 v1.6.0 及更高版本上对 PyTorch 可用。

如何使用优化器状态分片

您可以通过在 modelparallel 配置设置 "shard_optimizer_state": True 来启用优化器状态分片。

启用此功能后，库会根据数据并行度对模型参数集进行分区。与第 i 个分区对应的梯度，仅在第 i 个数据并行秩处缩减。在对 smp.step 修饰器函数的第一次调用结束时，被 smp.DistributedOptimizer 包装的优化器重新定义了其参数，使其仅限于与当前数据并行秩的分区相对应的参数。重新定义的参数称为虚拟参数，它们与原始参数共享底层存储。在第一次调用 optimizer.step 期间，优化器状态是根据这些重新定义的参数创建的，这些状态由于原始分区而被分片。优化器更新后，AllGather 操作（作为 optimizer.step 调用的一部分）跨数据并行秩运行，以实现一致的参数状态。

提示

当数据并行度大于 1 且模型的参数超过 10 亿时，优化器状态分片可能很有用。

数据并行性由 (processes_per_host * instance_count / pipeline_parallel_degree) 计算得出，smp.dp_size() 函数在后台处理大小调整。

配置 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,
        "shard_optimizer_state": True
    }
}

调整 PyTorch 训练脚本

请参阅将张量并行性于管道并行性结合使用部分中的调整 PyTorch 训练脚本。无需对脚本进行其他修改。