Entrenamiento con el FP16 con paralelismo de modelos - Amazon SageMaker AI

Las traducciones son generadas a través de traducción automática. En caso de conflicto entre la traducción y la version original de inglés, prevalecerá la version en inglés.

Entrenamiento con el FP16 con paralelismo de modelos

Para el entrenamiento del FP16, aplique las siguientes modificaciones al script y al estimador de entrenamiento.

nota

Esta característica está disponible para PyTorch en la biblioteca de paralelismo de modelos de SageMaker v1.10.0 y versiones posteriores.

Adaptar su script de entrenamiento de PyTorch

  1. Encapsule su modelo con el administrador de contexto 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 = ...
    sugerencia

    Si utiliza el paralelismo de tensores, agregue tensor_parallelism=smp.tp_size() > 1 al gestor contextual smp.model_creation. Agregar esta línea también ayuda a detectar automáticamente si el paralelismo de tensores está activado o no.

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

  2. Al encapsular el optimizador con smdistributed.modelparallel.torch.DistributedOptimizer, establezca el argumento static_loss_scaling o dynamic_loss_scaling. De forma predeterminada, static_loss_scaling se establece en 1.0 y dynamic_loss_scaling se establece en False. Si establece dynamic_loss_scale=True, puede introducir las opciones de escalado dinámico de pérdidas como un diccionario a través del argumento dynamic_loss_args. En la mayoría de los casos, le recomendamos utilizar el escalado de pérdidas dinámico con las opciones predeterminadas. Para obtener más información, opciones y ejemplos de la característica encapsuladora del optimizador, consulte la API smdistributed.modelparallel.torch.DistributedOptimizer.

    El siguiente código es un ejemplo de cómo encapsular un objeto optimizador Adadelta con una escala de pérdidas dinámica para el entrenamiento de 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
    }
)

Configurar un estimador PyTorch de SageMaker

Añada el parámetro FP16 ("fp16") a la configuración de distribución para el paralelismo del modelo al crear un objeto estimador PyTorch de SageMaker. Para obtener una lista completa de los parámetros de configuración del paralelismo de modelos, consulte Parámetros de 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(...)

Cuando se inicia el entrenamiento de FP16, el modelo y el optimizador están encapsulados por FP16_Module y FP16_Optimizer respectivamente, que son versiones smdistributed modificadas de las utilidades de Apex. FP16_Module convierte el modelo a FP16 dtype y se ocupa de la pasada hacia adelante en FP16.

sugerencia

Puede aplicar recorte de gradiente llamando clip_master_grads antes de optimizer.step.

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

Al utilizar torch.optim.lr_scheduler y un entrenamiento de FP16, tendrá que pasar optimizer.optimizer al programador LR en lugar de al optimizador. Vea el siguiente código de ejemplo:

from torch.optim.lr_scheduler import StepLR

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