使用模型并行性的 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
)