Addestramento FP16 con parallelismo dei modelli

Per l'addestramento sull'FP16, applica le seguenti modifiche allo script di addestramento e allo strumento di valutazione.

Nota

Questa funzionalità è disponibile PyTorch nella libreria di parallelismo dei SageMaker modelli v1.10.0 e versioni successive.

PyTorch Adatta il tuo script di allenamento

1. Effettua il wrapping del modello usando il gestore di contesto 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 = ...

   Suggerimento

   Se stai usando il parallelismo tensoriale, aggiungi tensor_parallelism=smp.tp_size() > 1 al gestore di contesto smp.model_creation. L'aggiunta di questa linea aiuta anche a rilevare automaticamente se il parallelismo tensoriale è attivato o meno.

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

2. Quando racchiudi l'ottimizzatore con smdistributed.modelparallel.torch.DistributedOptimizer, imposta l'argomento static_loss_scaling o dynamic_loss_scaling. Per impostazione predefinita, static_loss_scaling è impostato su 1.0 e dynamic_loss_scaling è impostato su False. Se imposti dynamic_loss_scale=True, puoi inserire le opzioni di dimensionamento dinamico delle perdite come dizionario tramite l'argomento dynamic_loss_args. Nella maggior parte dei casi, si consiglia di utilizzare la scala dinamica delle perdite con le opzioni predefinite. Per ulteriori informazioni, opzioni ed esempi della funzione wrapper dell'ottimizzatore, consultate smdistributed.modelparallel.torch. DistributedOptimizerAPI.

   Il codice seguente è un esempio di wrapping di un oggetto ottimizzatore Adadelta con una scala dinamica delle perdite per l'addestramento 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
       }
   )

Configura uno SageMaker PyTorch stimatore

Aggiungi il parametro FP16 ("fp16") alla configurazione di distribuzione per il parallelismo del modello durante la creazione di un oggetto estimatore. SageMaker PyTorch Per un elenco completo dei parametri di configurazione per il parallelismo dei modelli, vedere Parametri per 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(...)

Quando inizia l'addestramento su FP16, il modello e l'ottimizzatore vengono integrati rispettivamente da FP16_Module e FP16_Optimizer, versioni modificate smdistributed delle utilità Apex. FP16_Module converte il modello in FP16 dtype e gestisce il passaggio in avanti in FP16.

Suggerimento

È possibile applicare il ritaglio del gradiente chiamando clip_master_grads prima di optimizer.step.

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

Suggerimento

Quando si utilizza torch.optim.lr_scheduler e l'addestramento FP16, è necessario passare optimizer.optimizer al pianificatore LR anziché all'ottimizzatore. Guarda il codice di esempio seguente.

from torch.optim.lr_scheduler import StepLR

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