本文属于机器翻译版本。若本译文内容与英语原文存在差异，则一律以英文原文为准。
要进行 FP16 训练，请对训练脚本和估算器进行以下修改。
注意
此功能可在 SageMaker 模型并行度库 v1.10.0 及更高版本 PyTorch 中使用。
调整您的 PyTorch 训练脚本
-
使用 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 = ... -
当您使用
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和包装，它们是 Apexsmdistributed版本。 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
) 