优化器状态分片 - Amazon SageMaker
如何使用优化器状态分片

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

优化器状态分片

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

注意

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

如何使用优化器状态分片

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

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

提示

当数据并行度大于 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 训练脚本

请参阅 Tensor 并行度与流水线并行度相结合部分中的调整 PyTorch 训练脚本。无需对脚本进行其他修改。