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

翻訳は機械翻訳により提供されています。提供された翻訳内容と英語版の間で齟齬、不一致または矛盾がある場合、英語版が優先します。

モデル並列処理による 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.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 utils の修正 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 トレーニングを使用する場合は、オプティマイザではなく 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
)