La biblioteca paralela de modelos de Amazon SageMaker ahora acelera las cargas de trabajo de PyTorch FSDP hasta en un 20 %

La capacitación en modelos de lenguaje grande (LLM) ha ganado popularidad durante el último año con el lanzamiento de varios modelos populares como Llama 2, Falcon y Mistral. Los clientes ahora están preentrenando y ajustando LLM que van desde mil millones hasta más de 175 mil millones de parámetros para optimizar el rendimiento del modelo para aplicaciones en todas las industrias, desde atención médica hasta finanzas y marketing.

Formar modelos de alto rendimiento a esta escala puede ser un desafío. Los LLM de alta precisión pueden requerir terabytes de datos de entrenamiento y miles o incluso millones de horas de tiempo de cálculo del acelerador para lograr la precisión objetivo. Para completar la capacitación y lanzar productos de manera oportuna, los clientes confían en técnicas de paralelismo para distribuir esta enorme carga de trabajo entre hasta miles de dispositivos aceleradores. Sin embargo, estas técnicas de paralelismo pueden ser difíciles de usar: diferentes técnicas y bibliotecas solo son compatibles con ciertas cargas de trabajo o están restringidas a ciertas arquitecturas de modelo, el rendimiento del entrenamiento puede ser muy sensible a configuraciones oscuras y el estado del arte está evolucionando rápidamente. Como resultado, los profesionales del aprendizaje automático deben dedicar semanas de preparación para escalar sus cargas de trabajo de LLM a grandes grupos de GPU.

En esta publicación, destacamos las nuevas características de la biblioteca de modelos paralelos (SMP) de Amazon SageMaker que simplifican el proceso de capacitación de modelos grandes y lo ayudan a capacitar a los LLM más rápido. En particular, cubrimos la nueva experiencia de usuario simplificada de la biblioteca SMP que se basa en las API de datos paralelos totalmente fragmentados (FSDP) de PyTorch de código abierto, la funcionalidad paralela de tensor ampliada que permite entrenar modelos con cientos de miles de millones de parámetros y optimizaciones de rendimiento que reducen el tiempo de entrenamiento del modelo. y el costo hasta en un 20%.

Para obtener más información sobre la biblioteca paralela de modelos de SageMaker, consulte la documentación de la biblioteca de paralelismo de modelos de SageMaker v2. También puedes consultar nuestro cuadernos de ejemplo Para empezar.

Nuevas funciones que simplifican y aceleran el entrenamiento de modelos grandes

Esta publicación analiza las últimas funciones incluidas en la versión v2.0 de la biblioteca paralela de modelos SageMaker. Estas características mejoran la usabilidad de la biblioteca, amplían la funcionalidad y aceleran la capacitación. En las siguientes secciones, resumimos las nuevas funciones y analizamos cómo puede utilizar la biblioteca para acelerar el entrenamiento de su modelo grande.

Alineación de SMP con PyTorch de código abierto

Desde su lanzamiento en 2020, SMP ha permitido la capacitación a gran escala y de alto rendimiento en instancias informáticas de SageMaker. Con esta última versión principal de SMP, la biblioteca simplifica la experiencia del usuario al alinear sus API con PyTorch de código abierto.

Ofertas de PyTorch Paralelismo de datos totalmente fragmentados (FSDP) como su método principal para soportar grandes cargas de trabajo de capacitación en muchos dispositivos informáticos. Como se demuestra en el siguiente fragmento de código, las API actualizadas de SMP para técnicas como el paralelismo de datos fragmentados reflejan las de PyTorch. Puedes simplemente ejecutar import torch.sagemaker y usarlo en lugar de torch.

## training_script.py
import torch.sagemaker as tsm
tsm.init()

# Set up a PyTorch model
model = ...

# Wrap the PyTorch model using the PyTorch FSDP module
model = FSDP(
model,
...
)

optimizer = ...
...

Con estas actualizaciones de las API de SMP, ahora puede obtener los beneficios de rendimiento de SageMaker y la biblioteca SMP sin necesidad de revisar sus scripts de entrenamiento de PyTorch FSDP existentes. Este paradigma también le permite utilizar la misma base de código cuando entrena en las instalaciones que en SageMaker, simplificando la experiencia del usuario para los clientes que entrenan en múltiples entornos.

Para obtener más información sobre cómo habilitar SMP con sus scripts de entrenamiento de PyTorch FSDP existentes, consulte Introducción a SMP.

Integración del paralelismo tensorial para permitir el entrenamiento en clústeres masivos

Esta versión de SMP también amplía las capacidades de PyTorch FSDP para incluir técnicas de paralelismo tensorial. Un problema con el uso exclusivo del paralelismo de datos fragmentados es que puede encontrar problemas de convergencia a medida que aumenta el tamaño del clúster. Esto se debe a que la fragmentación de los parámetros, los gradientes y el estado del optimizador en rangos paralelos de datos también aumenta el tamaño del lote global; en grupos grandes, este tamaño de lote global se puede llevar más allá del umbral por debajo del cual convergería el modelo. Debe incorporar una técnica de paralelismo adicional que no requiera un aumento en el tamaño del lote global a medida que escala su clúster.

Para mitigar este problema, SMP v2.0 introduce la capacidad de componer paralelismo de datos fragmentados con paralelismo tensorial. El paralelismo tensorial permite que el tamaño del clúster aumente sin cambiar el tamaño del lote global ni afectar la convergencia del modelo. Con esta característica, puede aumentar de forma segura el rendimiento del entrenamiento aprovisionando clústeres con 256 nodos o más.

Hoy en día, el paralelismo tensorial con PyTorch FSDP solo está disponible con SMP v2. SMP v2 le permite habilitar esta técnica con unas pocas líneas de cambio de código y desbloquear un entrenamiento estable incluso en clústeres grandes. SMP v2 se integra con Motor transformador por su implementación del paralelismo tensorial y lo hace compatible con las API FSDP de PyTorch. Puede habilitar el paralelismo tensor de PyTorch FSDP y SMP simultáneamente sin realizar ningún cambio en su modelo de PyTorch o en la configuración de PyTorch FSDP. Los siguientes fragmentos de código muestran cómo configurar el diccionario de configuración SMP en formato JSON y agregar el módulo de inicialización SMP. torch.sagemaker.init()que acepta el diccionario de configuración en el backend cuando inicia el trabajo de capacitación, en su script de capacitación.

La configuración de SMP es la siguiente:

{
"tensor_parallel_degree": 8,
"tensor_parallel_seed": 0
}

En su script de entrenamiento, use el siguiente código:

import torch.sagemaker as tsm
tsm.init()

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_config(..)
model = tsm.transform(model)

Para obtener más información sobre el uso del paralelismo tensorial en SMP, consulte la sección sobre paralelismo tensorial de nuestra documentación.

Utilice funciones avanzadas para acelerar el entrenamiento de modelos hasta en un 20 %

Además de permitir la capacitación distribuida en clústeres con cientos de instancias, SMP también ofrece técnicas de optimización que pueden acelerar la capacitación de modelos hasta en un 20 %. En esta sección, destacamos algunas de estas optimizaciones. Para obtener más información, consulte la sección de funciones principales de nuestra documentación.

fragmentación híbrida

El paralelismo de datos fragmentados es una técnica de entrenamiento distribuido que ahorra memoria y divide el estado de un modelo (parámetros del modelo, gradientes y estados del optimizador) entre dispositivos. Esta huella de memoria más pequeña le permite incluir un modelo más grande en su clúster o aumentar el tamaño del lote. Sin embargo, el paralelismo de datos fragmentados también aumenta los requisitos de comunicación de su trabajo de capacitación porque los artefactos del modelo fragmentado con frecuencia se recopilan de diferentes dispositivos durante el entrenamiento. De esta manera, el grado de fragmentación es una configuración importante que compensa el consumo de memoria y la sobrecarga de comunicación.

De forma predeterminada, PyTorch FSDP fragmenta los artefactos del modelo en todos los dispositivos aceleradores de su clúster. Dependiendo de su trabajo de capacitación, este método de fragmentación podría aumentar los gastos generales de comunicación y crear un cuello de botella. Para ayudar con esto, la biblioteca SMP ofrece paralelismo de datos fragmentados híbridos configurables además de PyTorch FSDP. Esta función le permite establecer el grado de fragmentación óptimo para su carga de trabajo de entrenamiento. Simplemente especifique el grado de fragmentación en un objeto JSON de configuración e inclúyalo en su script de entrenamiento SMP.

La configuración de SMP es la siguiente:

{ "hybrid_shard_degree": 16 }

Para obtener más información sobre las ventajas del paralelismo de datos fragmentados híbridos, consulte Escalado casi lineal de entrenamiento de modelos gigantes en AWS. Para obtener más información sobre la implementación de fragmentación híbrida con su script de entrenamiento FSDP existente, consulte Paralelismo de datos compartidos híbridos en nuestra documentación.

Utilice las operaciones de comunicación colectiva SMDDP optimizadas para la infraestructura de AWS

Puede utilizar la biblioteca SMP con la biblioteca de paralelismo de datos distribuidos (SMDDP) de SageMaker para acelerar sus cargas de trabajo de capacitación distribuidas. SMDDP incluye un optimizado AllGather Operación de comunicación colectiva diseñada para obtener el mejor rendimiento en instancias aceleradas de SageMaker p4d y p4de. En la capacitación distribuida, las operaciones de comunicación colectiva se utilizan para sincronizar información entre los trabajadores de GPU. AllGather es una de las operaciones de comunicación colectiva principales que normalmente se utilizan en el paralelismo de datos fragmentados para materializar los parámetros de la capa antes de los pasos de cálculo hacia adelante y hacia atrás. Para los trabajos de capacitación que se ven obstaculizados por la comunicación, operaciones colectivas más rápidas pueden reducir el tiempo y el costo de la capacitación sin efectos secundarios sobre la convergencia.

Para utilizar la biblioteca SMDDP, solo necesita agregar dos líneas de código a su script de capacitación:

import torch.distributed as dist

# Initialize with SMDDP
import smdistributed.dataparallel.torch.torch_smddp
dist.init_process_group(backend="smddp") # Replacing "nccl"

# Initialize with SMP
import torch.sagemaker as tsm
tsm.init()

Además de SMP, SMDDP admite PyTorch FSDP y DeepSpeed ​​de código abierto. Para obtener más información sobre la biblioteca SMDDP, consulte Ejecutar capacitación distribuida con la biblioteca de paralelismo de datos distribuidos de SageMaker.

Descarga de activación

Normalmente, el paso hacia adelante del entrenamiento del modelo calcula las activaciones en cada capa y las mantiene en la memoria de la GPU hasta que finaliza el paso hacia atrás para la capa correspondiente. Estas activaciones almacenadas pueden consumir una cantidad significativa de memoria de la GPU durante el entrenamiento. La descarga de activación es una técnica que, en cambio, mueve estos tensores a la memoria de la CPU después del paso hacia adelante y luego los recupera a la GPU cuando son necesarios. Este enfoque puede reducir sustancialmente el uso de memoria de la GPU durante el entrenamiento.

Aunque PyTorch admite la descarga de activaciones, su implementación es ineficiente y puede hacer que las GPU estén inactivas mientras las activaciones se recuperan de la CPU durante un paso hacia atrás. Esto puede causar una degradación significativa del rendimiento cuando se utiliza la descarga de activación.

SMP v2 ofrece un algoritmo de descarga de activación optimizado que puede mejorar el rendimiento del entrenamiento. La implementación de SMP busca previamente las activaciones antes de que sean necesarias en la GPU, lo que reduce el tiempo de inactividad.

Debido a que SMP se basa en las API de PyTorch, habilitar la descarga de activación optimizada requiere solo unas pocas líneas de cambio de código. Simplemente agregue las configuraciones asociadas (sm_activation_offloading y activation_loading_horizon parámetros) e inclúyalos en su guión de entrenamiento.

La configuración de SMP es la siguiente:

{
"activation_loading_horizon": 2,
"sm_activation_offloading": True
}

En el script de entrenamiento, use el siguiente código:

import torch.sagemaker as tsm
tsm.init()

# Native PyTorch module for activation offloading
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
apply_activation_checkpointing,
offload_wrapper,
)

model = FSDP(...)

# Activation offloading requires activation checkpointing.
apply_activation_checkpointing(
model,
check_fn=checkpoint_tformer_layers_policy,
)

model = offload_wrapper(model)

Para obtener más información sobre las herramientas de punto de control de PyTorch de código abierto para la descarga de activación, consulte la checkpoint_wrapper.py script en el repositorio de PyTorch GitHub y Puntos de control de activación en la publicación del blog de PyTorch Escalado de modelos de cimentación multimodal en TorchMultimodal con Pytorch Distributed. Para obtener más información sobre la implementación optimizada de la descarga de activación de SMP, consulte la sección de descarga de activación de nuestra documentación.

Más allá de la fragmentación híbrida, SMDDP y la descarga de activación, SMP ofrece optimizaciones adicionales que pueden acelerar su gran carga de trabajo de entrenamiento de modelos. Esto incluye puntos de control de activación optimizados, inicialización retrasada de parámetros y otros. Para obtener más información, consulte la sección de funciones principales de nuestra documentación.

Conclusión

A medida que los conjuntos de datos, los tamaños de los modelos y los grupos de capacitación continúan creciendo, la capacitación distribuida eficiente es cada vez más crítica para la entrega oportuna y asequible de modelos y productos. La última versión de la biblioteca paralela de modelos SageMaker lo ayuda a lograr esto al reducir el cambio de código y alinearse con las API FSDP de PyTorch, lo que permite el entrenamiento en clústeres masivos mediante paralelismo tensorial y optimizaciones que pueden reducir el tiempo de entrenamiento hasta en un 20 %.

Para comenzar con SMP v2, consulte nuestra documentación y nuestra cuadernos de muestra.

Sobre los autores

Robert Van Dusen es gerente senior de productos en Amazon SageMaker. Lidera marcos, compiladores y técnicas de optimización para la capacitación en aprendizaje profundo.

luis quintela es el administrador de desarrolladores de software para la biblioteca paralela de modelos de AWS SageMaker. En su tiempo libre, se le puede encontrar montando su Harley en el Área de la Bahía de San Francisco.

Gautama Kumar es ingeniero de software con AWS ai Deep Learning. Le apasiona crear herramientas y sistemas para IA. En su tiempo libre, le gusta andar en bicicleta y leer libros.

Rahul Huilgol es ingeniero senior de desarrollo de software en aprendizaje profundo distribuido en Amazon Web Services.