FP16 モデル並列処理によるトレーニング

FP16 トレーニングでは、トレーニングスクリプトと推定器に次の変更を適用します。

注記

この機能は、 SageMaker モデル並列処理ライブラリ v1.10.0 以降 PyTorch で で使用できます。

PyTorch トレーニングスクリプトを適応させる

1. Wrap your model using the 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.torchDistributedOptimizer」を参照してくださいAPI。

   次のコードは、オ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_Moduleと でラップされます。これらは FP16_Optimizer Apex utils の変更smdistributedバージョンです。 はモデルを FP16 dtype FP16_Moduleに変換し、 のフォワードパスを処理しますFP16。 https://nvidia.github.io/apex/fp16_utils.html#apex-fp16-utils

ヒント

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トレーニングを使用する場合は、オプティマイザではなく LR スケジューラoptimizer.optimizerに を渡す必要があります。次のコード例を参照してください。

from torch.optim.lr_scheduler import StepLR

scheduler = StepLR(
    optimizer.optimizer if smp.state.cfg.fp16 else optimizer,
    step_size=1,
    gamma=args.gamma
)