使用模型并行性的 FP16 训练 - Amazon SageMaker
Amazon Web Services 文档中描述的 Amazon Web Services 服务或功能可能因区域而异。要查看适用于中国区域的差异，请参阅 中国的 Amazon Web Services 服务入门 (PDF)。

使用模型并行性的 FP16 训练

对于 FP16 训练，请对训练脚本和估算器应用以下修改。

注意

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

调整 PyTorch 训练脚本

  1. 使用 smdistributed.modelparallel.torch.model_creation() 上下文管理器包装您的模型。

    # fp16_training_script.py

import torch
import smdistributed.modelparallel.torch as smp

with smp.model_creation(
    dtype=torch.float16 if args.fp16 else torch.get_default_dtype()
):
    model = ...
    提示

    如果您使用的是张量并行性，请将 tensor_parallelism=smp.tp_size() > 1 添加到 smp.model_creation 上下文管理器中。添加此行还有助于自动检测是否已激活张量并行性。

    with smp.model_creation(
    ... ,
    tensor_parallelism=smp.tp_size() > 1
):
    model = ...

  2. 当您使用 smdistributed.modelparallel.torch.DistributedOptimizer 包装优化器时，请设置 static_loss_scaling 或 dynamic_loss_scaling 参数。默认情况下，static_loss_scaling 设置为 1.0，dynamic_loss_scaling 设置为 False。如果您设置 dynamic_loss_scale=True，则可以通过 dynamic_loss_args 参数将动态损失缩放选项作为字典输入。在大多数情况下，我们建议您使用带有默认选项的动态损失缩放。有关优化器包装器函数的更多信息、选项和示例，请参阅 smdistributed.modelparallel.torch.DistributedOptimizer API。

    以下代码是在 FP16 训练中使用动态损失缩放包装 Adadelta 优化器对象的示例。

    optimizer = torch.optim.Adadelta(...)
optimizer = smp.DistributedOptimizer(
    optimizer,
    static_loss_scale=None,
    dynamic_loss_scale=True,
    dynamic_loss_args={
        "scale_window": 1000,
        "min_scale": 1,
        "delayed_shift": 2
    }
)

配置 SageMaker PyTorch 估算器

在创建 SageMaker PyTorch 估算器对象时，将 FP16 参数 ("fp16") 添加到分布配置中以实现模型并行性。有关模型并行性配置参数的完整列表，请参阅 smdistributed 的参数。

from sagemaker.pytorch import PyTorch

smp_options = {
    "enabled": True,
    "parameters":  {
        "microbatches":  4,
        "pipeline_parallel_degree":  2,
        "tensor_parallel_degree":  2,
        ...,

        "fp16": True
    }
}

fp16_estimator = PyTorch(
    entry_point="fp16_training_script.py", # Specify your train script
    ...,

    distribution={
        "smdistributed": {"modelparallel": smp_options},
        "mpi": {...}
    }
)

fp16_estimator.fit(...)

当 FP16 训练开始时，模型和优化器分别由 FP16_Module 和 FP16_Optimizer 包装，这是 Apex 实用程序修改后的 smdistributed 版本。FP16_Module 将模型转换为 FP16 dtype 并处理 FP16 中的向前传递。

提示

您可以在 optimizer.step 之前通过调用 clip_master_grads 来应用梯度剪裁。

optimizer.clip_master_grads(max_norm)     # max_norm(float or int): max norm of the gradients
提示

使用 torch.optim.lr_scheduler 和 FP16 训练时，您需要将 optimizer.optimizer 传递给 LR 计划程序而不是优化器。请参阅以下示例代码。

from torch.optim.lr_scheduler import StepLR

scheduler = StepLR(
    optimizer.optimizer if smp.state.cfg.fp16 else optimizer,
    step_size=1,
    gamma=args.gamma
)