Acelera PyTorch con DeepSpeed para entrenar modelos de lenguaje grandes con instancias DL1 basadas en Intel Habana Gaudi en EC2.

Speed up PyTorch training of large language models using DeepSpeed with DL1 instances based on Intel Habana Gaudi on EC2.

Entrenar modelos de lenguaje grandes (LLMs) con miles de millones de parámetros puede ser desafiante. Además de diseñar la arquitectura del modelo, los investigadores necesitan establecer técnicas de entrenamiento de vanguardia para el entrenamiento distribuido como el soporte de precisión mixta, la acumulación de gradientes y la creación de checkpoints. Con modelos grandes, la configuración de entrenamiento es aún más desafiante porque la memoria disponible en un solo dispositivo acelerador limita el tamaño de los modelos entrenados utilizando solo paralelismo de datos, y el uso del entrenamiento paralelo del modelo requiere un nivel adicional de modificaciones en el código de entrenamiento. Bibliotecas como DeepSpeed (una biblioteca de optimización de aprendizaje profundo de código abierto para PyTorch) abordan algunos de estos desafíos y pueden ayudar a acelerar el desarrollo y entrenamiento del modelo.

En esta publicación, configuramos el entrenamiento en las instancias de Amazon Elastic Compute Cloud (Amazon EC2) DL1 basadas en Intel Habana Gaudi y cuantificamos los beneficios de usar un marco de escalado como DeepSpeed. Presentamos resultados de escalado para un modelo transformador tipo codificador (BERT con 340 millones a 1.5 mil millones de parámetros). Para el modelo de 1.5 mil millones de parámetros, logramos una eficiencia de escalado del 82.7% en 128 aceleradores (16 instancias dl1.24xlarge) usando las optimizaciones de DeepSpeed ZeRO stage 1. Los estados del optimizador fueron particionados por DeepSpeed para entrenar grandes modelos utilizando el paradigma de paralelismo de datos. Este enfoque se ha extendido para entrenar un modelo de 5 mil millones de parámetros utilizando paralelismo de datos. También usamos el soporte nativo de Gaudi del tipo de datos BF16 para reducir el tamaño de la memoria y aumentar el rendimiento de entrenamiento en comparación con el tipo de datos FP32. Como resultado, logramos la convergencia del modelo de pre-entrenamiento (fase 1) en 16 horas (nuestro objetivo era entrenar un modelo grande en un día) para el modelo BERT de 1.5 mil millones de parámetros utilizando el conjunto de datos wikicorpus-en.

Configuración de entrenamiento

Provisionamos un clúster de computación administrado compuesto por 16 instancias dl1.24xlarge utilizando AWS Batch. Desarrollamos un taller de AWS Batch que ilustra los pasos para configurar el clúster de entrenamiento distribuido con AWS Batch. Cada instancia dl1.24xlarge tiene ocho aceleradores Habana Gaudi, cada uno con 32 GB de memoria y una red RoCE de malla completa entre tarjetas con un ancho de banda de interconexión bidireccional total de 700 Gbps cada uno (consulte Amazon EC2 DL1 instances Deep Dive para obtener más información). El clúster dl1.24xlarge también utilizó cuatro adaptadores de tela elástica de AWS (EFA), con un total de 400 Gbps de interconexión entre nodos.

El taller de entrenamiento distribuido ilustra los pasos para configurar el clúster de entrenamiento distribuido. El taller muestra la configuración de entrenamiento distribuido utilizando AWS Batch y, en particular, la función de trabajos paralelos de múltiples nodos para lanzar trabajos de entrenamiento contenerizados a gran escala en clústeres completamente administrados. Más específicamente, se crea un entorno de cómputo AWS Batch completamente administrado con instancias DL1. Los contenedores se extraen del Registro de contenedores elástico de Amazon (Amazon ECR) y se lanzan automáticamente en las instancias del clúster en función de la definición de trabajo paralelo de múltiples nodos. El taller concluye ejecutando un entrenamiento de paralelismo de datos de múltiples nodos y múltiples HPU de un modelo BERT (de 340 millones a 1.5 mil millones de parámetros) utilizando PyTorch y DeepSpeed.

Pre-entrenamiento de BERT 1.5B con DeepSpeed

Las versiones 1.5 y 1.6 de Habana SynapseAI admiten las optimizaciones de DeepSpeed ZeRO1. El fork de Habana del repositorio de GitHub de DeepSpeed incluye las modificaciones necesarias para admitir los aceleradores Gaudi. Hay soporte completo de paralelismo de datos distribuido (multi-tarjeta, multi-instancia), optimizaciones ZeRO1 y tipos de datos BF16.

• Presentamos la sintonización de popularidad para artículos similare...
• Los gigantes de la electrónica se adentran en la automatización ind...
• NVIDIA lleva la autonomía avanzada a los robots móviles con Isaac AMR.

Todas estas características están habilitadas en el repositorio de referencia del modelo BERT 1.5B, que presenta un modelo de codificador bidireccional de 48 capas, 1600 dimensiones ocultas y 25 cabezas, derivado de una implementación de BERT. El repositorio también contiene la implementación del modelo BERT Large de referencia: una arquitectura de red neuronal de 24 capas, 1024 dimensiones ocultas, 16 cabezas y 340 millones de parámetros. Los scripts de modelado de pre-entrenamiento se derivan del repositorio de Ejemplos de aprendizaje profundo de NVIDIA para descargar los datos wikicorpus_en, preprocesar los datos en bruto en tokens y dividir los datos en conjuntos de datos h5 más pequeños para el entrenamiento de paralelismo de datos distribuido. Puede adoptar este enfoque genérico para entrenar sus arquitecturas de modelo PyTorch personalizadas utilizando sus conjuntos de datos utilizando instancias DL1.

Resultados de escalado de pre-entrenamiento (fase 1)

Para el pre-entrenamiento de modelos grandes a escala, nos enfocamos principalmente en dos aspectos de la solución: el rendimiento de entrenamiento, medido por el tiempo de entrenamiento, y la rentabilidad de llegar a una solución completamente convergente. A continuación, profundizamos en estas dos métricas con el pre-entrenamiento de BERT 1.5B como ejemplo.

Escalabilidad de rendimiento y tiempo de entrenamiento

Comenzamos midiendo el rendimiento de la implementación de BERT Large como línea de base para la escalabilidad. La siguiente tabla enumera el rendimiento medido de secuencias por segundo de 1-8 instancias dl1.24xlarge (con ocho dispositivos aceleradores por instancia). Utilizando la eficiencia de rendimiento de una sola instancia como línea de base, medimos la eficiencia de la escalabilidad a través de múltiples instancias, lo cual es una palanca importante para entender la métrica de entrenamiento de relación precio-rendimiento.

La siguiente figura ilustra la eficiencia de escalabilidad.

Para BERT 1.5B, modificamos los hiperparámetros del modelo en el repositorio de referencia para garantizar la convergencia. El tamaño efectivo del lote por acelerador se estableció en 384 (para la utilización máxima de memoria), con micro lotes de 16 por paso y 24 pasos de acumulación de gradiente. Se utilizaron tasas de aprendizaje de 0,0015 y 0,003 para 8 y 16 nodos, respectivamente. Con estas configuraciones, logramos la convergencia de la fase 1 de pre-entrenamiento de BERT 1.5B en 8 instancias dl1.24xlarge (64 aceleradores) en aproximadamente 25 horas y en 15 horas en 16 instancias dl1.24xlarge (128 aceleradores). La siguiente figura muestra la pérdida promedio en función del número de épocas de entrenamiento, a medida que aumentamos el número de aceleradores.

Con la configuración descrita anteriormente, obtuvimos una eficiencia de escalabilidad fuerte del 85% con 64 aceleradores y del 83% con 128 aceleradores, a partir de una línea de base de 8 aceleradores en una sola instancia. La siguiente tabla resume los parámetros.

La siguiente figura ilustra la eficiencia de escalado.

Conclusión

En este post, evaluamos el soporte para DeepSpeed por Habana SynapseAI v1.5/v1.6 y cómo ayuda a escalar el entrenamiento LLM en los aceleradores Habana Gaudi. La pre-entrenamiento de un modelo BERT de 1.5 mil millones de parámetros tardó 16 horas en converger en un clúster de 128 aceleradores Gaudi, con una escalabilidad fuerte del 85%. Les animamos a echar un vistazo a la arquitectura demostrada en el taller de AWS y considerar adoptarla para entrenar arquitecturas de modelos PyTorch personalizados utilizando instancias DL1.

• Advanced (300)
• Amazon EC2