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

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

分片优化器状态

优化程序状态分片是一种有用的节省内存的技术，它可以跨数据并行设备组对优化程序状态（描述优化程序状态的权重集）进行分片。无论何时使用有状态优化器（例如亚当）或 FP16 优化器（存储参数的 FP16 和 FP32 副本），都可以使用优化程序状态分片。

如何使用优化器状态分片

优化程序状态分片可以通过设置打开"shard_optimizer_state": True中的modelparallel配置。

启用此功能后，库将根据数据并行度对模型参数集进行分区。对应于i'第个分区只能在i'第个数据 parallel 排名。在第一次通话结束时smp.step装饰器函数，包装的优化器smp.DistributedOptimizer将其参数重新定义为仅限于与当前数据 parallel 排名的分区相对应的参数。重新定义的参数被称为virtual_parameters并使用原始参数共享底层存储。在第一次打电话期间optimizer.step，优化程序状态是基于这些重新定义的参数创建的，这些参数是由于原始分区而进行分片的。优化程序更新后，allgather操作（作为optimizer.step调用) 在数据并行等级中执行，以实现一致的参数状态。

提示

当数据并行度大于 1 且模型具有超过十亿个参数时，优化程序状态分片可能非常有用。

数据并行度的计算方法是(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 训练脚本中的Tensor 并行度与管道并行度相结合部分。脚本不需要进行额外的修改。