FP16使用模型并行度进行训练 - Amazon SageMaker

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

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
)