PyTorch - Amazon SageMaker AI
Modelos de TensorFlow con transformadores Hugging Face

PyTorch

Lleve su propio modelo de PyTorch a SageMaker AI y ejecute el trabajo de entrenamiento con el compilador de entrenamiento de SageMaker.

Modelos de TensorFlow con transformadores Hugging Face

Los modelos de PyTorch con Hugging Face Transformers se basan en la API torch.nn.Module de PyTorch. Hugging Face Transformers también ofrece clases de modelos de Trainer y previamente entrenadas para PyTorch para ayudar a reducir el esfuerzo de configurar modelos de procesamiento de lenguaje natural (NLP). Después de preparar el script de entrenamiento, puede lanzar un trabajo de entrenamiento usando el estimador PyTorch o HuggingFace de SageMaker AI con la configuración del Compilador de entrenamiento de SageMaker cuando pase al siguiente tema en Habilitar Compilador de entrenamiento de SageMaker.

sugerencia

Cuando cree un tokenizador para un modelo de PNL con Transformers en su script de entrenamiento, asegúrese de utilizar una forma de tensor de entrada estática especificando padding='max_length'. No utilice padding='longest' porque rellenar la secuencia más larga del lote puede cambiar la forma del tensor de cada lote de entrenamiento. La forma de entrada dinámica puede desencadenar la recompilación del modelo y aumentar el tiempo total de entrenamiento. Para obtener más información sobre las opciones de relleno de los tokenizadores de Transformers, consulte Padding and truncation en la documentación de Hugging Face Transformers.

Grandes modelos lingüísticos mediante la clase Trainer de Hugging Face Transformers

Si utiliza la clase Trainer de la biblioteca de transformadores, no necesitará realizar ningún cambio adicional en su script de entrenamiento. el Compilador de entrenamiento de SageMaker compila automáticamente su modelo Trainer si lo activa a través de la clase estimador. El siguiente código muestra la forma básica de un script de entrenamiento PyTorch con Hugging Face Trainer API.

from transformers import Trainer, TrainingArguments

training_args=TrainingArguments(**kwargs)
trainer=Trainer(args=training_args, **kwargs)

Para el entrenamiento con una sola GPU

No necesita cambiar su código cuando utilice la clase transformers.Trainer.

Para entrenamiento distribuido

PyTorch v1.11.0 y superior

Para ejecutar el entrenamiento distribuido con el Compilador de Entrenamiento SageMaker, debe agregar la siguiente función _mp_fn() en su script de entrenamiento y encapsular la función main(). Redirige las llamadas a función _mp_fn(index) desde el tiempo de ejecución distribuido de SageMaker AI para PyTorch (pytorchxla) a la función main() del script de entrenamiento. 

def _mp_fn(index):
    main()

Esta función acepta el argumento index para indicar el rango de la GPU actual en el clúster para el entrenamiento distribuido. Para encontrar más scripts de ejemplo, consulte los scripts de ejemplo de modelado lingüístico de Hugging Face Transformers.

Para Transformers v4.17 y anteriores con PyTorch v1.10.2 y anteriores

El compilador de entrenamiento de SageMaker utiliza un mecanismo alternativo para lanzar un trabajo de entrenamiento distribuido, y no es necesario realizar ninguna modificación en el script de entrenamiento. En su lugar, el Compilador de entrenamiento de SageMaker requiere que pase un script lanzador de entrenamiento distribuido de SageMaker AI al argumento entry_point y que pase su script de entrenamiento al argumento hyperparameters del estimador Hugging Face de SageMaker AI.

Prácticas recomendadas para utilizar el Compilador de entrenamiento de SageMaker con Trainer

  • Asegúrese de que utiliza optimizadores SyncFree estableciendo el argumento optim en adamw_torch_xla al configurar transformers.TrainingArgument. Consulte también Optimizer en la documentación de Hugging Face Transformers.

  • Asegúrese de que el rendimiento del canal de procesamiento de datos es superior al rendimiento de entrenamiento. Puede modificar los argumentos dataloader_num_workers y preprocessing_num_workers de la clase transformers.TrainingArgument para conseguirlo. Normalmente, deben ser mayores o iguales que el número de GPU, pero menores que el número de CPU.

Cuando haya terminado de adaptar su script de entrenamiento, continúe con Ejecutar trabajos de entrenamiento de PyTorch con el compilador de entrenamiento de SageMaker.

Grandes modelos lingüísticos utilizando PyTorch directamente (sin la API Hugging Face Transformers Trainer)

Si tiene un script de entrenamiento que utiliza PyTorch directamente, deberá realizar cambios adicionales en su script de entrenamiento de PyTorch para implementar PyTorch/XLA. Siga las instrucciones para modificar su script para configurar correctamente los primitivos de PyTorch/XLA.

Para el entrenamiento con una sola GPU

  1. Importe las bibliotecas de optimización.

    import torch_xla
import torch_xla.core.xla_model as xm

  2. Cambiar el dispositivo de destino para que sea XLA en lugar de torch.device("cuda") 

    device=xm.xla_device()

  3. Si está utilizando Automatic Mixed Precision (AMP) de PyTorch, haga lo siguiente:

    1. Reemplace torch.cuda.amp por lo siguiente:

      import torch_xla.amp

    2. Reemplace torch.optim.SGD y torch.optim.Adam por lo siguiente:

      import torch_xla.amp.syncfree.Adam as adam
import torch_xla.amp.syncfree.SGD as SGD

    3. Reemplace torch.cuda.amp.GradScaler por lo siguiente:

      import torch_xla.amp.GradScaler as grad_scaler

  4. Si no utiliza AMP, reemplace optimizer.step() por lo siguiente:

    xm.optimizer_step(optimizer)

  5. Si estás usando un cargador de datos distribuido, encapsule su dataloader en la clase ParallelLoader de PyTorch/XLA:

    import torch_xla.distributed.parallel_loader as pl
parallel_loader=pl.ParallelLoader(dataloader, [device]).per_device_loader(device)

  6. Añada mark_step al final del bucle de entrenamiento cuando no estés utilizando parallel_loader:

    xm.mark_step()

  7. Para comprobar su entrenamiento, utilice el método de comprobación de modelos de PyTorch/XLA:

    xm.save(model.state_dict(), path_to_save)

Cuando haya terminado de adaptar su script de entrenamiento, continúe con Ejecutar trabajos de entrenamiento de PyTorch con el compilador de entrenamiento de SageMaker.

Para entrenamiento distribuido

Además de los cambios enumerados en la sección Para el entrenamiento con una sola GPU anterior, añada los siguientes cambios para distribuir correctamente la carga de trabajo entre las GPU.

  1. Si esta utilizando AMP, añada all_reduce después de scaler.scale(loss).backward():

    gradients=xm._fetch_gradients(optimizer)
xm.all_reduce('sum', gradients, scale=1.0/xm.xrt_world_size())

  2. Si necesita establecer variables para local_ranks y world_size, utilice un código similar al siguiente:

    local_rank=xm.get_local_ordinal()
world_size=xm.xrt_world_size()

  3. Para cualquier world_size (num_gpus_per_node*num_nodes) mayor que 1, debe definir un muestreador de entrenamiento que debe tener un aspecto similar al siguiente:

    import torch_xla.core.xla_model as xm

if xm.xrt_world_size() > 1:
    train_sampler=torch.utils.data.distributed.DistributedSampler(
        train_dataset,
        num_replicas=xm.xrt_world_size(),
        rank=xm.get_ordinal(),
        shuffle=True
    )

train_loader=torch.utils.data.DataLoader(
    train_dataset, 
    batch_size=args.batch_size,
    sampler=train_sampler,
    drop_last=args.drop_last,
    shuffle=False if train_sampler else True,
    num_workers=args.num_workers
)

  4. Realice los siguientes cambios para asegurarse de que utiliza el parallel_loader proporcionado por el módulo torch_xla distributed. 

    import torch_xla.distributed.parallel_loader as pl
train_device_loader=pl.MpDeviceLoader(train_loader, device)

    El train_device_loader funciona como un cargador PyTorch normal de la siguiente manera: 

    for step, (data, target) in enumerate(train_device_loader):
    optimizer.zero_grad()
    output=model(data)
    loss=torch.nn.NLLLoss(output, target)
    loss.backward()

    Con todos estos cambios, debería ser capaz de lanzar un entrenamiento distribuido con cualquier modelo PyTorch sin la API Transformer Trainer. Tenga en cuenta que estas instrucciones se pueden utilizar tanto para varias GPU de un solo nodo como para varias GPU de varios nodos.

  5. Para PyTorch v1.11.0 y posterior

    Para ejecutar el entrenamiento distribuido con el Compilador de Entrenamiento SageMaker, debe agregar la siguiente función _mp_fn() en su script de entrenamiento y encapsular la función main(). Redirige las llamadas a función _mp_fn(index) desde el tiempo de ejecución distribuido de SageMaker AI para PyTorch (pytorchxla) a la función main() del script de entrenamiento. 

    def _mp_fn(index):
    main()

    Esta función acepta el argumento index para indicar el rango de la GPU actual en el clúster para el entrenamiento distribuido. Para encontrar más scripts de ejemplo, consulte los scripts de ejemplo de modelado lingüístico de Hugging Face Transformers.

    Para Transformers v4.17 y anteriores con PyTorch v1.10.2 y anteriores

    El Compilador de entrenamiento de SageMaker utiliza un mecanismo alternativo para lanzar un trabajo de entrenamiento distribuido y requiere que pase un script lanzador de entrenamiento distribuido de SageMaker AI al argumento entry_point y que pase su script de entrenamiento al argumento hyperparameters del estimador Hugging Face de SageMaker AI.

Cuando haya terminado de adaptar su script de entrenamiento, continúe con Ejecutar trabajos de entrenamiento de PyTorch con el compilador de entrenamiento de SageMaker.

Prácticas recomendadas para utilizar el Compilador de entrenamiento de SageMaker con PyTorch/XLA

Si desea aprovechar el Compilador de entrenamiento de SageMaker en su script de entrenamiento nativo de PyTorch, es posible que desee familiarizarse primero con PyTorch en dispositivos XLA. En las siguientes secciones se enumeran algunas de las prácticas recomendadas para habilitar XLA para PyTorch.

nota

En esta sección de prácticas recomendadas se asume que utiliza los siguientes módulos de PyTorch/XLA:

import torch_xla.core.xla_model as xm
import torch_xla.distributed.parallel_loader as pl
Comprender el modo Lazy en PyTorch/XLA

Una diferencia significativa entre PyTorch/XLA y PyTorch nativo es que el sistema PyTorch/XLA se ejecuta en modo Lazy mientras que el PyTorch nativo se ejecuta en modo Eager. Los tensores en modo Lazy son marcadores de posición para construir el gráfico computacional hasta que se materializan después de que la compilación y la evaluación se hayan completado. El sistema PyTorch/XLA construye el gráfico computacional sobre la marcha cuando llama a las API de PyTorch para generar el cálculo usando tensores y operadores. El gráfico computacional se compila y ejecuta cuando xm.mark_step() es llamado explícita o implícitamente por pl.MpDeviceLoader/pl.ParallelLoader, o cuando se solicita explícitamente el valor de un tensor como por ejemplo llamando a loss.item() o print(loss).

Minimizar el número de compilation-and-executions utilizando pl.MpDeviceLoader/pl.ParallelLoader y xm.step_closure

Para obtener el mejor rendimiento, debe tener en cuenta las posibles formas de iniciar compilación y ejecuciones como se describe en Comprender el modo Lazy en PyTorch/XLA e intentar minimizar el número de compilación y ejecuciones. Idealmente, solo es necesaria una compilación y ejecución por iteración de entrenamiento y se inicia automáticamente mediante pl.MpDeviceLoader/pl.ParallelLoader. El MpDeviceLoader está optimizado para XLA y debe utilizarse siempre que sea posible para obtener el mejor rendimiento. Durante el entrenamiento, es posible que desee examinar algunos resultados intermedios, como los valores de pérdida. En tal caso, la impresión de tensores Lazy debe encapsularse utilizando xm.add_step_closure() para evitar compilación y ejecuciones innecesarias.

Utilizar AMP y optimizadores syncfree

El entrenamiento en modo Automatic Mixed Precision (AMP) acelera considerablemente la velocidad de entrenamiento aprovechando los kernels tensoriales de las GPU NVIDIA. El Compilador de entrenamiento de SageMaker proporciona optimizadores syncfree que están optimizados para XLA para mejorar el rendimiento de AMP. Actualmente, están disponibles los tres optimizadores syncfree siguientes, que deben utilizarse si es posible para obtener el mejor rendimiento.

torch_xla.amp.syncfree.SGD
torch_xla.amp.syncfree.Adam
torch_xla.amp.syncfree.AdamW

Estos optimizadores syncfree deben emparejarse con torch_xla.amp.GradScaler para el escalado/desescalado de gradientes.

sugerencia

A partir de PyTorch 1.13.1, el Compilador de entrenamiento de Sagemaker mejora el rendimiento al permitir que PyTorch/XLA anule automáticamente los optimizadores (como SGD, Adam, AdamW) en torch.optim o transformers.optimization con sus versiones sin sincronización en torch_xla.amp.syncfree (como torch_xla.amp.syncfree.SGD, torch_xla.amp.syncfree.Adam, torch_xla.amp.syncfree.AdamW). No tiene que cambiar las líneas de código en las que define los optimizadores en su script de entrenamiento.