FP16 使用模型并行度进行训练 - 亚马逊 SageMaker AI
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。 DistributedOptimizerAPI。

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

    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_Optimizer分别由FP16_Module和包装，它们是 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
)