

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

# 优化器状态分片
<a name="model-parallel-extended-features-pytorch-optimizer-state-sharding"></a>

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

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

## 如何使用优化器状态分片
<a name="model-parallel-extended-features-pytorch-optimizer-state-sharding-how-to-use"></a>

您可以通过在 `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 训练脚本](model-parallel-extended-features-pytorch-tensor-parallelism-examples.md#model-parallel-extended-features-pytorch-tensor-and-pipeline-parallelism-script)。无需对脚本进行其他修改。