FP16-Training mit Modellparallelität

Nehmen Sie für das FP16-Training die folgenden Änderungen an Ihrem Trainingsskript und Ihrem Schätzer vor.

Anmerkung

Diese Funktion ist PyTorch in der SageMaker Modellparallelismus-Bibliothek v1.10.0 und höher verfügbar.

Passen Sie Ihr Trainingsskript an PyTorch

1. Umschließen Sie Ihr Modell mit dem Kontextmanager 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 = ...

   Tipp

   Wenn Sie Tensor-Parallelität verwenden, fügen Sie tensor_parallelism=smp.tp_size() > 1 zum smp.model_creationKontextmanager hinzu. Mit Hilfe dieser zusätzlichen Zeile kann auch automatisch erkannt werden, ob die Tensor-Parallelität aktiviert ist oder nicht.

   with smp.model_creation(
       ... ,
       tensor_parallelism=smp.tp_size() > 1
   ):
       model = ...

2. Wenn Sie den Optimierer mit smdistributed.modelparallel.torch.DistributedOptimizer umschließen, setzen Sie entweder das Argument static_loss_scaling oder dynamic_loss_scaling. static_loss_scaling Ist standardmäßig auf 1.0 gesetzt und dynamic_loss_scaling ist auf False gesetzt. Wenn Sie dynamic_loss_scale=True einstellen, können Sie dynamische Verlustskalierungsoptionen als Wörterbuch über das Argument dynamic_loss_args einspeisen. In den meisten Fällen empfehlen wir die dynamische Verlustskalierung mit den Standardoptionen zu verwenden. Weitere Informationen, Optionen und Beispiele für die Optimizer-Wrapper-Funktion finden Sie unter smdistributed.modelparallel.torch. DistributedOptimizerAPI.

   Der folgende Code ist ein Beispiel für das Umschließen eines Adadelta Optimizer-Objekts mit dynamischer Verlustskalierung für das FP16-Training.

   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
       }
   )

Konfigurieren Sie einen SageMaker PyTorch Schätzer

Fügen Sie der Verteilungskonfiguration den FP16-Parameter ("fp16") für Modellparallelität hinzu, wenn Sie ein Schätzerobjekt erstellen. SageMaker PyTorch Eine vollständige Liste der Konfigurationsparameter für Modellparallelität finden Sie unter Parameter für 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(...)

Wenn das FP16-Training beginnt, werden Modell und Optimizer von FP16_Module und bzw. FP16_Optimizer umschlossen. Dabei handelt es sich um modifizierte smdistributed Versionen der Apex-Hilfsprogramme. FP16_Module konvertiert das Modell in FP16 dtype und kümmert sich um den Vorwärtsdurchgang in FP16.

Tipp

Sie können die Steigungen beschneiden, indem Sie clip_master_grads vor optimizer.step aufrufen.

optimizer.clip_master_grads(max_norm)     # max_norm(float or int): max norm of the gradients

Tipp

Wenn Sie das torch.optim.lr_scheduler FP16-Training verwenden, müssen Sie optimizer.optimizer an den LR-Scheduler übergeben, und nicht an den Optimizer. Schauen Sie sich den folgenden Beispiel-Code an:

from torch.optim.lr_scheduler import StepLR

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