FP16使用模型平行度進行訓練 - Amazon SageMaker

本文為英文版的機器翻譯版本，如內容有任何歧義或不一致之處，概以英文版為準。

FP16使用模型平行度進行訓練

若要進行FP16訓練，請將下列修改套用至訓練指令碼和估算器。

注意

此功能適用於 PyTorch SageMaker 模型平行程度庫 v1.10.0 及更新版本。

調整您的 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 引數饋送動態損失縮放選項作為字典。在多數情況，我們建議您使用動態損失縮放與預設選項。如需最佳化程式包裝函式的詳細資訊、選項和範例，請參閱小分散式 .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
)