La Tecnología Detrás del Entrenamiento BLOOM

La tecnología de entrenamiento BLOOM.

En los últimos años, se ha vuelto habitual entrenar modelos de lenguaje cada vez más grandes. Si bien se discute con frecuencia el tema de que esos modelos no se publiquen para su estudio posterior, se presta muy poca atención al conocimiento oculto sobre cómo entrenar dichos modelos. Este artículo tiene como objetivo cambiar esto al arrojar algo de luz sobre la tecnología e ingeniería detrás del entrenamiento de dichos modelos, tanto en términos de hardware como de software, utilizando como ejemplo el modelo de lenguaje BLOOM con 176 mil millones de parámetros.

Pero primero nos gustaría agradecer a las empresas, personas clave y grupos que hicieron posible la increíble hazaña de entrenar un modelo de 176 mil millones de parámetros con un pequeño grupo de personas dedicadas.

Luego se discutirá la configuración de hardware y los componentes tecnológicos principales.

• Ventaja Actor Critic (A2C)
• Presentamos nueva documentación de audio y visión en 🤗 Datasets
• Nyströmformer Aproximando la autoatención en tiempo lineal y memoria utilizando el método de Nyström.

Aquí hay un resumen rápido del proyecto:

Personas

El proyecto fue concebido por Thomas Wolf (cofundador y CSO – Hugging Face), quien se atrevió a competir con las grandes corporaciones no solo para entrenar uno de los modelos multilingües más grandes, sino también para hacer que el resultado final sea accesible para todas las personas, convirtiendo así lo que era solo un sueño para la mayoría de las personas en realidad.

Este artículo se centra específicamente en el aspecto de ingeniería del entrenamiento del modelo. La parte más importante de la tecnología detrás de BLOOM fueron las personas y las empresas que compartieron su experiencia y nos ayudaron con la codificación y el entrenamiento.

Hay 6 grupos principales de personas a quienes agradecer:

1. El equipo BigScience de HuggingFace, que dedicó más de media docena de empleados a tiempo completo para descubrir y ejecutar el entrenamiento desde el inicio hasta la meta final, y proporcionó e pagó toda la infraestructura más allá de la computación de Jean Zay.
2. El equipo Microsoft DeepSpeed, que desarrolló DeepSpeed y luego lo integró con Megatron-LM, y cuyos desarrolladores pasaron muchas semanas trabajando en las necesidades del proyecto y proporcionando muchos consejos prácticos y experienciales asombrosos antes y durante el entrenamiento.
3. El equipo NVIDIA Megatron-LM, que desarrolló Megatron-LM y que fue de gran ayuda al responder nuestras numerosas preguntas y brindar consejos de primera clase basados en la experiencia.
4. El equipo IDRIS / GENCI que administra la supercomputadora Jean Zay, que donó al proyecto una cantidad increíble de potencia de cálculo y un excelente soporte de administración del sistema.
5. El equipo PyTorch que creó un marco súper potente, en el que se basó el resto del software, y que nos brindó un gran apoyo durante la preparación para el entrenamiento, corrigiendo múltiples errores y mejorando la usabilidad de los componentes de PyTorch en los que nos basamos durante el entrenamiento.
6. Los voluntarios en el grupo de trabajo de ingeniería de BigScience

Sería muy difícil nombrar a todas las personas increíbles que contribuyeron al aspecto de ingeniería del proyecto, así que solo mencionaré a algunas personas clave fuera de Hugging Face que fueron el fundamento de ingeniería de este proyecto durante los últimos 14 meses:

Olatunji Ruwase, Deepak Narayanan, Jeff Rasley, Jared Casper, Samyam Rajbhandari y Rémi Lacroix

También estamos agradecidos con todas las empresas que permitieron que sus empleados contribuyeran a este proyecto.

Descripción general

La arquitectura de BLOOM es muy similar a la de GPT3 con algunas mejoras adicionales que se discutirán más adelante en este artículo.

El modelo se entrenó en Jean Zay, la supercomputadora financiada por el gobierno francés que es administrada por GENCI y se encuentra instalada en IDRIS, el centro nacional de computación del Centro Nacional de Investigación Científica (CNRS) de Francia. La potencia de cálculo fue generosamente donada al proyecto por GENCI (subvención 2021-A0101012475).

Se utilizó el siguiente hardware durante el entrenamiento:

• GPUs: 384 NVIDIA A100 80GB GPUs (48 nodos) + 32 GPUs de repuesto
• 8 GPUs por nodo utilizando NVLink 4 para la conexión entre GPUs, 4 enlaces OmniPath
• CPU: Procesador AMD EPYC 7543 de 32 núcleos
• Memoria de la CPU: 512GB por nodo
• Memoria de la GPU: 640GB por nodo
• Conexión entre nodos: Arquitectura Omni-Path (OPA) con árbol gordiano sin bloqueo
• Red de comunicaciones NCCL: una subred totalmente dedicada
• Red de E/S de disco: GPFS compartido con otros nodos y usuarios

Puntos de control:

• puntos de control principales
• cada punto de control con estados de optimización fp32 y pesos de bf16+fp32 tiene un tamaño de 2.3TB – solo los pesos de bf16 tienen un tamaño de 329GB.

Conjuntos de datos:

• 46 idiomas en 1.5TB de texto masivo deduplicado y limpiado, convertido en 350B tokens únicos
• El tamaño del vocabulario del modelo es de 250,680 tokens
• Para más detalles, consulte The BigScience Corpus A 1.6TB Composite Multilingual Dataset

El entrenamiento del modelo BLOOM de 176B ocurrió entre marzo y julio de 2022 y tardó aproximadamente 3.5 meses en completarse (aproximadamente 1 millón de horas de cálculo).

Megatron-DeepSpeed

El modelo BLOOM de 176B ha sido entrenado utilizando Megatron-DeepSpeed, que es una combinación de 2 tecnologías principales:

• DeepSpeed es una biblioteca de optimización de aprendizaje profundo que facilita, eficiente y efectivo el entrenamiento distribuido.
• Megatron-LM es un marco de modelo de transformador grande y poderoso desarrollado por el equipo de Investigación de Aprendizaje Profundo Aplicado de NVIDIA.

El equipo de DeepSpeed desarrolló una implementación basada en paralelismo 3D combinando la fragmentación ZeRO y el paralelismo de canalización de la biblioteca DeepSpeed con el paralelismo de tensor de Megatron-LM. Se pueden ver más detalles sobre cada componente en la tabla a continuación.

Tenga en cuenta que el Megatron-DeepSpeed de BigScience es un fork del repositorio original de Megatron-DeepSpeed, al cual hemos agregado múltiples adiciones.

A continuación se muestra una tabla que indica qué componentes fueron proporcionados por qué marco para entrenar BLOOM:

Tenga en cuenta que tanto Megatron-LM como DeepSpeed tienen implementaciones de paralelismo de canalización y optimizador BF16, pero usamos los de DeepSpeed ya que están integrados con ZeRO.

Megatron-DeepSpeed implementa el paralelismo 3D para permitir que los modelos enormes se entrenen de manera muy eficiente. Veamos brevemente los componentes 3D.

1. DataParallel (DP) – la misma configuración se replica varias veces, y cada una recibe una porción de los datos. El procesamiento se realiza en paralelo y todas las configuraciones se sincronizan al final de cada paso de entrenamiento.
2. TensorParallel (TP) – cada tensor se divide en múltiples fragmentos, por lo que en lugar de que el tensor completo resida en una sola GPU, cada fragmento del tensor reside en su GPU designada. Durante el procesamiento, cada fragmento se procesa por separado y en paralelo en diferentes GPU y los resultados se sincronizan al final del paso. Esto es lo que se podría llamar paralelismo horizontal, ya que la división ocurre a nivel horizontal.
3. PipelineParallel (PP) – el modelo se divide verticalmente (a nivel de capa) en múltiples GPU, de modo que solo una o varias capas del modelo se colocan en una sola GPU. Cada GPU procesa en paralelo diferentes etapas del canal y trabaja en un pequeño fragmento del lote.
4. Zero Redundancy Optimizer (ZeRO) – también realiza la fragmentación de los tensores de manera algo similar a TP, excepto que el tensor completo se reconstruye a tiempo para un cálculo hacia adelante o hacia atrás, por lo que el modelo no necesita ser modificado. También admite diversas técnicas de desplazamiento para compensar la memoria limitada de la GPU.

Paralelismo de datos

La mayoría de los usuarios con solo unas pocas GPU probablemente estén familiarizados con DistributedDataParallel (DDP) en la documentación de PyTorch. En este método, el modelo se replica completamente en cada GPU y luego, después de cada iteración, todos los modelos sincronizan sus estados entre sí. Este enfoque permite acelerar el entrenamiento al asignar más recursos al problema, pero solo funciona si el modelo cabe en una sola GPU.

Paralelismo de datos ZeRO

El paralelismo de datos ZeRO (ZeRO-DP) se describe en el siguiente diagrama de esta publicación de blog

Puede ser difícil entenderlo, pero en realidad, el concepto es bastante simple. Esto es solo el DDP habitual, excepto que, en lugar de replicar los parámetros completos del modelo, los gradientes y los estados del optimizador, cada GPU almacena solo una parte de ellos. Y luego, en tiempo de ejecución, cuando se necesitan los parámetros completos de la capa para la capa dada, todas las GPU se sincronizan para proporcionarse mutuamente las partes que les faltan, eso es todo.

Este componente está implementado por DeepSpeed.

Paralelismo de tensores

En el paralelismo de tensores (TP), cada GPU procesa solo una porción de un tensor y solo agrega el tensor completo para operaciones que requieren el tensor completo.

En esta sección utilizamos conceptos y diagramas del artículo de Megatron-LM: Entrenamiento eficiente de modelos de lenguaje a gran escala en clústeres de GPU.

El elemento principal de cualquier transformador es una capa completamente conectada nn.Linear seguida de una activación no lineal GeLU.

Siguiendo la notación del artículo de Megatron, podemos escribir la parte del producto punto como Y = GeLU(XA), donde X e Y son los vectores de entrada y salida, y A es la matriz de pesos.

Si observamos la computación en forma de matriz, es fácil ver cómo la multiplicación de matrices se puede dividir entre múltiples GPUs:

Si dividimos la matriz de pesos A columna por columna entre N GPUs y realizamos las multiplicaciones de matrices XA_1 a través de XA_n en paralelo, entonces obtendremos N vectores de salida Y_1, Y_2, ..., Y_n que se pueden alimentar a GeLU de forma independiente: . Observa que con la matriz Y dividida a lo largo de las columnas, podemos dividir la segunda GEMM a lo largo de sus filas para que tome la salida de GeLU directamente sin ninguna comunicación adicional.

Utilizando este principio, podemos actualizar una MLP de profundidad arbitraria, mientras sincronizamos las GPUs después de cada secuencia de filas-columnas. Los autores del artículo de Megatron-LM proporcionan una ilustración útil para eso:

Aquí, f es un operador de identidad en el paso hacia adelante y una reducción total en el paso hacia atrás, mientras que g es una reducción total en el paso hacia adelante y una identidad en el paso hacia atrás.

La paralelización de las capas de atención con múltiples cabezas es aún más simple, ¡ya que son inherentemente paralelas debido a tener múltiples cabezas independientes!

Consideraciones especiales: Debido a las dos reducciones totales por capa tanto en los pasos hacia adelante como hacia atrás, TP requiere una interconexión muy rápida entre dispositivos. Por lo tanto, no es recomendable hacer TP en más de un nodo, a menos que tengas una red muy rápida. En nuestro caso, la interconexión entre nodos era mucho más lenta que PCIe. En la práctica, si un nodo tiene 4 GPUs, el grado más alto de TP es 4. Si necesitas un grado de TP de 8, debes usar nodos que tengan al menos 8 GPUs.

Este componente está implementado por Megatron-LM. Megatron-LM ha ampliado recientemente el paralelismo a nivel de tensor para incluir el paralelismo de secuencia que divide las operaciones que no se pueden dividir como se mencionó anteriormente, como LayerNorm, a lo largo de la dimensión de secuencia. El artículo Reducing Activation Recomputation in Large Transformer Models proporciona detalles sobre esta técnica. El paralelismo de secuencia se desarrolló después de que BLOOM fuera entrenado, por lo que no se utilizó en el entrenamiento de BLOOM.

Paralelismo en Pipelines

El Paralelismo en Pipelines Naive (PP ingenuo) es donde se distribuyen grupos de capas del modelo entre múltiples GPUs y simplemente se mueve los datos de una GPU a otra como si fuera una única GPU compuesta. El mecanismo es relativamente simple: cambiar las capas deseadas a los dispositivos deseados con .to() y ahora, cuando los datos ingresan y salen de esas capas, cambiar los datos al mismo dispositivo que la capa y dejar el resto sin modificar.

Esto realiza un paralelismo vertical del modelo, porque si recuerdas cómo se dibujan la mayoría de los modelos, dividimos las capas verticalmente. Por ejemplo, si el siguiente diagrama muestra un modelo de 8 capas:

===================  ===================
|  0 | 1 | 2 | 3  |  |  4 | 5 | 6 | 7  |
===================  ===================
        GPU0                 GPU1

simplemente lo dividimos en 2 verticalmente, colocando las capas 0-3 en la GPU0 y las capas 4-7 en la GPU1.

Ahora, mientras los datos viajan de la capa 0 a la 1, de la 1 a la 2 y de la 2 a la 3, esto es similar al paso hacia adelante de un modelo normal en una sola GPU. Pero cuando los datos necesitan pasar de la capa 3 a la capa 4, necesitan viajar de la GPU0 a la GPU1, lo que introduce una sobrecarga de comunicación. Si las GPUs participantes están en el mismo nodo de cómputo (por ejemplo, la misma máquina física), esta copia es bastante rápida, pero si las GPUs están ubicadas en diferentes nodos de cómputo (por ejemplo, varias máquinas), la sobrecarga de comunicación podría ser significativamente mayor.

Luego, las capas 4 a 5 a 6 a 7 son como las tendría un modelo normal y cuando la séptima capa se completa, a menudo necesitamos enviar los datos de vuelta a la capa 0 donde están las etiquetas (o alternativamente enviar las etiquetas a la última capa). Ahora se puede calcular la pérdida y el optimizador puede hacer su trabajo.

Problemas:

• la deficiencia principal y por qué se llama PP “ingenuo” es que todos menos una GPU están inactivas en un momento dado. Entonces, si se utilizan 4 GPUs, es casi idéntico a cuadruplicar la cantidad de memoria de una sola GPU y se ignora el resto del hardware. Además, está el sobrecoste de copiar los datos entre dispositivos. Entonces, 4 tarjetas de 6 GB podrán acomodar el mismo tamaño que una tarjeta de 24 GB utilizando PP ingenuo, excepto que esta última completará el entrenamiento más rápido, ya que no tiene el sobrecoste de copiar los datos. Pero, supongamos que tiene tarjetas de 40 GB y necesita adaptar un modelo de 45 GB, puede hacerlo con 4 tarjetas de 40 GB (pero apenas debido al gradiente y los estados del optimizador).
• las incrustaciones compartidas pueden necesitar copiarse de ida y vuelta entre las GPUs.

La Paralelización en Tubería (PP) es casi idéntica a una PP ingenua descrita anteriormente, pero resuelve el problema de inactividad de la GPU dividiendo el lote entrante en micro lotes y creando artificialmente una tubería, que permite que diferentes GPUs participen simultáneamente en el proceso de cálculo.

La siguiente ilustración del documento GPipe muestra la PP ingenua en la parte superior y la PP en la parte inferior:

Es fácil ver en el diagrama inferior cómo la PP tiene menos zonas muertas, donde las GPUs están inactivas. Las partes inactivas se conocen como la “burbuja”.

Ambas partes del diagrama muestran un paralelismo de grado 4. Es decir, 4 GPUs participan en la tubería. Así que hay un camino directo de 4 etapas de tubería F0, F1, F2 y F3 y luego un camino inverso de regreso en orden inverso B3, B2, B1 y B0.

PP introduce un nuevo hiperparámetro que se puede ajustar, llamado chunks. Define cuántos fragmentos de datos se envían en secuencia a través de la misma etapa de tubería. Por ejemplo, en el diagrama inferior, se puede ver que chunks=4. GPU0 realiza el mismo camino directo en los fragmentos 0, 1, 2 y 3 (F0,0, F0,1, F0,2, F0,3) y luego espera a que otras GPUs hagan su trabajo y solo cuando su trabajo comienza a completarse, GPU0 vuelve a trabajar haciendo el camino inverso para los fragmentos 3, 2, 1 y 0 (B0,3, B0,2, B0,1, B0,0).

Observa que conceptualmente este es el mismo concepto que los pasos de acumulación de gradientes (GAS). PyTorch utiliza chunks, mientras que DeepSpeed se refiere al mismo hiperparámetro como GAS.

Debido a los fragmentos, PP introduce el concepto de micro lotes (MBS). DP divide el tamaño del lote de datos global en mini lotes, por lo que si tiene un grado DP de 4, un tamaño de lote global de 1024 se divide en 4 mini lotes de 256 cada uno (1024/4). Y si el número de chunks (o GAS) es 32, obtendremos un tamaño de micro lote de 8 (256/32). Cada etapa de tubería trabaja con un solo micro lote a la vez.

Para calcular el tamaño del lote global de la configuración DP + PP, hacemos: mbs*chunks*dp_degree (8*32*4=1024).

Volvamos al diagrama.

Con chunks=1 obtenemos la PP ingenua, que es muy ineficiente. Con un valor de chunks muy grande, obtenemos tamaños de micro lotes muy pequeños, lo que también podría no ser muy eficiente. Por lo tanto, es necesario experimentar para encontrar el valor que conduzca a la utilización más eficiente de las GPUs.

Aunque el diagrama muestra que hay una burbuja de tiempo “muerto” que no se puede paralelizar porque la última etapa de forward tiene que esperar a que se complete el backward de la tubería, el propósito de encontrar el mejor valor para chunks es permitir una alta utilización concurrente de las GPUs participantes, lo que se traduce en minimizar el tamaño de la burbuja.

Este mecanismo de programación se conoce como todos hacia adelante, todos hacia atrás. Algunas otras alternativas son uno hacia adelante, uno hacia atrás e intercalado uno hacia adelante, uno hacia atrás.

Aunque tanto Megatron-LM como DeepSpeed tienen su propia implementación del protocolo PP, Megatron-DeepSpeed utiliza la implementación de DeepSpeed ya que está integrada con otros aspectos de DeepSpeed.

Otro problema importante aquí es el tamaño de la matriz de incrustación de palabras. Normalmente, una matriz de incrustación de palabras consume menos memoria que el bloque transformador, pero en nuestro caso, con un vocabulario enorme de 250k, la capa de incrustación necesitaba 7.2GB en pesos bf16 y el bloque transformador solo 4.9GB. Por lo tanto, tuvimos que instruir a Megatron-Deepspeed para que considerara la capa de incrustación como un bloque transformador. Así que tuvimos una canalización de 72 capas, 2 de las cuales estaban dedicadas a la incrustación (la primera y la última). Esto permitió equilibrar el consumo de memoria de la GPU. Si no lo hubiéramos hecho, las etapas primera y última consumirían la mayor parte de la memoria de la GPU, y el 95% de las GPU usarían mucha menos memoria, por lo que el entrenamiento estaría lejos de ser eficiente.

DP+PP

El siguiente diagrama del tutorial de la canalización de DeepSpeed muestra cómo combinar DP con PP.

Aquí es importante ver cómo el rango DP 0 no ve la GPU2 y el rango DP 1 no ve la GPU3. Para DP solo existen las GPU 0 y 1, donde alimenta datos como si hubiera solo 2 GPU. GPU0 “secretamente” descarga parte de su carga en la GPU2 usando PP. Y GPU1 hace lo mismo con GPU3.

Dado que cada dimensión requiere al menos 2 GPU, aquí necesitarías al menos 4 GPU.

DP+PP+TP

Para obtener un entrenamiento aún más eficiente, PP se combina con TP y DP, lo que se denomina paralelismo 3D. Esto se puede ver en el siguiente diagrama.

Este diagrama es de una publicación del blog Paralelismo 3D: Escalando a modelos de billones de parámetros, que también es interesante de leer.

Dado que cada dimensión requiere al menos 2 GPU, aquí necesitarías al menos 8 GPU para un paralelismo 3D completo.

ZeRO DP+PP+TP

Una de las características principales de DeepSpeed es ZeRO, que es una extensión de DP súper escalable. Ya se ha discutido en ZeRO Data Parallelism. Normalmente es una función independiente que no requiere PP ni TP. Pero se puede combinar con PP y TP.

Cuando ZeRO-DP se combina con PP (y opcionalmente TP), generalmente solo habilita la etapa 1 de ZeRO, que divide solo los estados del optimizador. La etapa 2 de ZeRO divide adicionalmente los gradientes y la etapa 3 también divide los pesos del modelo.

Aunque teóricamente es posible utilizar la etapa 2 de ZeRO con el paralelismo de canalización, tendrá malos impactos en el rendimiento. Se necesitaría una colectiva adicional de reducción-dispersión para cada micro-batch para agregar los gradientes antes de la división, lo que añade una sobrecarga de comunicación potencialmente significativa. Debido a la naturaleza del paralelismo de canalización, se utilizan micro-batches pequeños y el enfoque se centra en tratar de equilibrar la intensidad aritmética (tamaño del micro-batch) con minimizar la burbuja de la canalización (número de micro-batches). Por lo tanto, esos costos de comunicación serán perjudiciales.

Además, ya hay menos capas que lo normal debido a PP, por lo que los ahorros de memoria no serán enormes. PP ya reduce el tamaño del gradiente en 1/PP, por lo que los ahorros de división de gradientes sobre eso son menos significativos que en DP puro.

La etapa 3 de ZeRO también se puede utilizar para entrenar modelos a esta escala, sin embargo, requiere más comunicación que la implementación paralela 3D de DeepSpeed. Después de una evaluación cuidadosa en nuestro entorno que ocurrió hace un año, encontramos que el paralelismo 3D de Megatron-DeepSpeed tenía el mejor rendimiento. Desde entonces, el rendimiento de la etapa 3 de ZeRO ha mejorado drásticamente y si lo evaluáramos hoy, quizás habríamos elegido la etapa 3 en su lugar.

BF16Optimizer

Entrenar modelos LLM enormes en FP16 no es recomendable.

Lo hemos demostrado a nosotros mismos al pasar varios meses entrenando un modelo 104B que, como puedes ver en el tensorboard, fue un completo fracaso. Aprendimos muchas cosas mientras luchábamos contra la divergencia constante de la pérdida de lm:

y también recibimos el mismo consejo de los equipos de Megatron-LM y DeepSpeed después de entrenar el modelo 530B. El reciente lanzamiento de OPT-175B también informó que tuvieron muchas dificultades para entrenar en FP16.

Así que en enero, sabiendo que íbamos a entrenar en A100s que admiten el formato BF16, Olatunji Ruwase desarrolló un BF16Optimizer que utilizamos para entrenar BLOOM.

Si no estás familiarizado con este formato de datos, echa un vistazo a la disposición de los bits. La clave del formato BF16 es que tiene el mismo exponente que FP32 y, por lo tanto, no sufre de desbordamiento como FP16 sufre mucho. Con FP16, que tiene un rango numérico máximo de 64k, solo puedes multiplicar números pequeños. Por ejemplo, puedes hacer 250*250=62500, pero si intentas 255*255=65025 te encontrarás con un desbordamiento, que es lo que causa los principales problemas durante el entrenamiento. Esto significa que tus pesos deben permanecer pequeños. Una técnica llamada escalamiento de pérdida puede ayudar con este problema, pero el rango limitado de FP16 sigue siendo un problema cuando los modelos son muy grandes.

BF16 no tiene ese problema, puedes hacer fácilmente 10_000*10_000=100_000_000 y no hay problema.

Por supuesto, dado que BF16 y FP16 tienen el mismo tamaño de 2 bytes, no se obtiene un almuerzo gratis y se paga con una precisión realmente mala al usar BF16. Sin embargo, si recuerdas, el entrenamiento utilizando descenso de gradiente estocástico y sus variaciones es algo así como una caminata titubeante, por lo que si no obtienes la dirección perfecta de inmediato, no hay problema, te corregirás en los siguientes pasos.

Independientemente de si se usa BF16 o FP16, también hay una copia de los pesos que siempre está en FP32, esto es lo que actualiza el optimizador. Por lo tanto, los formatos de 16 bits solo se utilizan para los cálculos, el optimizador actualiza los pesos FP32 con precisión completa y luego los convierte al formato de 16 bits para la siguiente iteración.

Todos los componentes de PyTorch se han actualizado para asegurarse de que realicen cualquier acumulación en FP32, por lo que no hay pérdida allí.

Un problema crucial es la acumulación de gradientes, y es una de las principales características del paralelismo de canalización, ya que los gradientes de cada procesamiento de microbatches se acumulan. Es crucial implementar la acumulación de gradientes en FP32 para mantener el entrenamiento preciso, y eso es lo que hace BF16Optimizer.

Además de otras mejoras, creemos que el entrenamiento de precisión mixta con BF16 convirtió una pesadilla potencial en un proceso relativamente fluido, como se puede observar en el siguiente gráfico de pérdida de lm:

Núcleos de CUDA fusionados

La GPU realiza dos cosas. Puede copiar datos hacia/desde la memoria y realizar cálculos en esos datos. Mientras la GPU está ocupada copiando, las unidades de cálculo de la GPU están inactivas. Si queremos utilizar eficientemente la GPU, queremos minimizar el tiempo inactivo.

Un kernel es un conjunto de instrucciones que implementa una operación específica de PyTorch. Por ejemplo, cuando llamas a torch.add, pasa por un despachador de PyTorch que examina el tensor o tensores de entrada y varias otras cosas, y decide qué código debe ejecutar y luego lo ejecuta. Un kernel de CUDA es una implementación específica que utiliza la biblioteca de API de CUDA y solo se puede ejecutar en GPUs NVIDIA.

Ahora, cuando le indicamos a la GPU que calcule c = torch.add(a, b); e = torch.max([c,d]), un enfoque ingenuo, y lo que PyTorch hará a menos que se indique lo contrario, es lanzar dos kernels separados, uno para realizar la suma de a y b y otro para encontrar el valor máximo entre c y d. En este caso, la GPU obtiene de su memoria a y b, realiza la suma y luego copia el resultado de vuelta en la memoria. Luego, obtiene c y d y realiza la operación max y nuevamente copia el resultado de vuelta en la memoria.

Si fusionáramos estas dos operaciones, es decir, las colocamos en un solo “kernel fusionado” y solo lanzamos ese kernel, no copiaremos el resultado intermedio c en la memoria, sino que lo dejaremos en los registros de la GPU y solo necesitaremos obtener d para completar el último cálculo. Esto ahorra mucha sobrecarga y evita la inactividad de la GPU, lo que hace que toda la operación sea mucho más eficiente.

Los kernels fusionados son precisamente eso. Principalmente reemplazan múltiples cálculos discretos y movimientos de datos hacia/desde la memoria por cálculos fusionados que tienen muy pocos movimientos de memoria. Además, algunos kernels fusionados reescriben las matemáticas para que ciertos grupos de cálculos se puedan realizar más rápido.

Para entrenar BLOOM de forma rápida y eficiente, fue necesario utilizar varios kernels CUDA fusionados personalizados proporcionados por Megatron-LM. En particular, hay un kernel optimizado para realizar LayerNorm, así como kernels para fusionar varias combinaciones de operaciones de escalado, enmascaramiento y softmax. La adición de un término de sesgo también se fusiona con la operación GeLU utilizando la funcionalidad JIT de PyTorch. Estas operaciones están limitadas por la memoria, por lo que es importante fusionarlas para maximizar la cantidad de cálculos realizados una vez que se ha recuperado un valor de la memoria. Por lo tanto, por ejemplo, agregar el término de sesgo mientras se realiza la operación GeLU limitada por la memoria no agrega tiempo adicional. Todos estos kernels están disponibles en el repositorio de Megatron-LM .

Conjuntos de datos

Otra característica importante de Megatron-LM es el cargador de datos eficiente. Durante el inicio del entrenamiento inicial, cada conjunto de datos se divide en muestras de la longitud de secuencia solicitada (2048 para BLOOM) y se crea un índice para enumerar cada muestra. Según los parámetros de entrenamiento, se calcula el número de épocas para un conjunto de datos y se crea un orden para esa cantidad de épocas y luego se baraja. Por ejemplo, si un conjunto de datos tiene 10 muestras y se debe recorrer dos veces, el sistema primero coloca los índices de las muestras en orden [0, ..., 9, 0, ..., 9] y luego baraja ese orden para crear el orden global final para el conjunto de datos. Observe que esto significa que el entrenamiento no simplemente recorrerá todo el conjunto de datos y luego se repetirá, es posible ver la misma muestra dos veces antes de ver otra muestra en absoluto, pero al final del entrenamiento el modelo habrá visto cada muestra dos veces. Esto ayuda a garantizar una curva de entrenamiento suave durante todo el proceso de entrenamiento. Estos índices, incluidos los desplazamientos en el conjunto de datos base de cada muestra, se guardan en un archivo para evitar volver a calcularlos cada vez que se inicia un proceso de entrenamiento. Varios de estos conjuntos de datos se pueden combinar con pesos variables en los datos finales vistos por el proceso de entrenamiento.

LayerNorm de incrustación

Mientras luchábamos por evitar que 104B se divergiera, descubrimos que agregar un LayerNorm adicional justo después de la primera incrustación de palabras hacía que el entrenamiento fuera mucho más estable.

Esta idea surgió de experimentar con bitsandbytes que contiene una StableEmbedding que es una incrustación normal con layernorm y utiliza una inicialización uniforme de Xavier.

Codificación posicional

También reemplazamos la incrustación posicional habitual con un AliBi, basado en el artículo: Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation, que permite extrapolar para secuencias de entrada más largas que las que se utilizaron para entrenar el modelo. Por lo tanto, aunque entrenamos con secuencias de longitud 2048, el modelo también puede manejar secuencias mucho más largas durante la inferencia.

Dificultades en el entrenamiento

Con la arquitectura, el hardware y el software en su lugar, pudimos comenzar a entrenar a principios de marzo de 2022. Sin embargo, no fue un proceso sin dificultades a partir de ahí. En esta sección discutimos algunos de los principales obstáculos que encontramos.

Había muchas cuestiones que resolver antes de que comenzara el entrenamiento. En particular, encontramos varios problemas que se manifestaron solo una vez que comenzamos a entrenar en 48 nodos y no aparecerían a pequeña escala. Por ejemplo, se necesitaba CUDA_LAUNCH_BLOCKING=1 para evitar que el marco se bloqueara, y tuvimos que dividir los grupos de optimizadores en grupos más pequeños; de lo contrario, el marco se bloquearía nuevamente. Puedes leer más detalles sobre esto en las crónicas previas al entrenamiento .

El principal tipo de problema encontrado durante el entrenamiento fueron las fallas de hardware. Como se trataba de un clúster nuevo con alrededor de 400 GPU, en promedio teníamos 1-2 fallas de GPU por semana. Estábamos guardando un punto de control cada 3 horas (100 iteraciones), por lo que en promedio perderíamos 1.5 horas de entrenamiento por cada bloqueo de hardware. Los administradores del sistema de Jean Zay reemplazaban las GPU defectuosas y volvían a poner en funcionamiento el nodo. Mientras tanto, teníamos nodos de respaldo para usar en su lugar.

Nos hemos encontrado con una variedad de otros problemas que causaron un tiempo de inactividad de 5-10 horas varias veces, algunos relacionados con un error de bloqueo en PyTorch, otros debido a la falta de espacio en disco. Si estás interesado en detalles específicos, consulta las crónicas de entrenamiento.

Estábamos planeando todos estos tiempos de inactividad al decidir sobre la viabilidad de entrenar este modelo: elegimos el tamaño del modelo para que coincidiera con esa viabilidad y la cantidad de datos que queríamos que el modelo consumiera. Con todos los tiempos de inactividad logramos terminar el entrenamiento en nuestro tiempo estimado. Como se mencionó anteriormente, llevó aproximadamente 1 millón de horas de cálculo completarlo.

Otro problema fue que SLURM no fue diseñado para ser utilizado por un equipo de personas. Un trabajo de SLURM es propiedad de un solo usuario y si no están presentes, los demás miembros del grupo no pueden hacer nada con el trabajo en ejecución. Desarrollamos una solución alternativa de apagado que permite a otros usuarios del grupo detener el proceso actual sin necesidad de que el usuario que inició el proceso esté presente. Esto funcionó bien en el 90% de los problemas. Si los diseñadores de SLURM leen esto, por favor agreguen el concepto de grupos de Unix para que un trabajo de SLURM pueda ser propiedad de un grupo.

Dado que el entrenamiento estaba ocurriendo las 24 horas del día, los 7 días de la semana, necesitábamos que alguien estuviera disponible en caso de emergencia, pero como teníamos personas tanto en Europa como en la costa oeste de Canadá, no era necesario que alguien llevara un localizador, simplemente nos superponíamos bien. Por supuesto, alguien tenía que supervisar el entrenamiento los fines de semana también. Automatizamos la mayoría de las cosas, incluida la recuperación de fallas de hardware, pero a veces también se necesitaba intervención humana.

Conclusión

La parte más difícil e intensa del entrenamiento fue el período de 2 meses previo al inicio del entrenamiento. Estábamos bajo mucha presión para comenzar el entrenamiento lo antes posible, ya que la asignación de recursos tenía un límite de tiempo y no tuvimos acceso a los A100 hasta el último momento. Así que fue un momento muy difícil, considerando que el BF16Optimizer se escribió en el último momento y necesitábamos depurarlo y corregir varios errores. Y como se explica en la sección anterior, descubrimos nuevos problemas que se manifestaron solo una vez que comenzamos a entrenar en 48 nodos y no aparecerán a pequeña escala.

Pero una vez que resolvimos esos problemas, el entrenamiento en sí fue sorprendentemente fluido y sin problemas importantes. La mayor parte del tiempo, una persona monitoreaba el entrenamiento y solo algunas veces varias personas participaban para solucionar problemas. Tuvimos un gran apoyo de la administración de Jean Zay, que rápidamente resolvió la mayoría de las necesidades que surgieron durante el entrenamiento.

En general, fue una experiencia muy intensa pero muy gratificante.

El entrenamiento de modelos de lenguaje grandes sigue siendo una tarea desafiante, pero esperamos que al construir y compartir esta tecnología de manera abierta, otros puedan basarse en nuestra experiencia.

Recursos

Enlaces importantes

• documento principal de entrenamiento
• tensorboard
• script de slurm de entrenamiento
• crónicas de entrenamiento

Artículos y documentos

No podríamos haber explicado todo en detalle en este artículo, así que si la tecnología presentada aquí ha despertado tu curiosidad y quieres saber más, aquí están los documentos para leer:

Megatron-LM:

• Entrenamiento eficiente de modelos de lenguaje a gran escala en clústeres de GPU.
• Reduciendo la recomputación de activaciones en modelos de transformador grandes.

DeepSpeed:

• ZeRO: Optimizaciones de memoria para entrenamiento de modelos de parámetros trillonarios.
• ZeRO-Offload: Democratizando el entrenamiento de modelos a escala de miles de millones.
• ZeRO-Infinity: Rompiendo la barrera de memoria de la GPU para el aprendizaje profundo a gran escala.
• DeepSpeed: Entrenamiento de modelos a escala extrema para todos.

Megatron-LM y Deepspeed juntos:

• Usando DeepSpeed y Megatron para entrenar a Megatron-Turing NLG 530B, un modelo de lenguaje generativo a gran escala.

ALiBi:

• Entrenar a corto plazo, probar a largo plazo: la atención con sesgos lineales permite la extrapolación de la longitud de entrada.
• ¿Qué modelo de lenguaje entrenar si tienes un millón de horas de GPU? – allí encontrarás los experimentos que nos llevaron a elegir ALiBi.

BitsNBytes:

• Optimizadores de 8 bits a través de cuantización por bloques (en el contexto de Embedding LayerNorm, pero el resto del artículo y la tecnología son increíbles, la única razón por la que no estábamos usando el optimizador de 8 bits es porque ya estábamos ahorrando memoria del optimizador con DeepSpeed-ZeRO).

Créditos del blog

Un gran agradecimiento a las siguientes personas amables que hicieron buenas preguntas y ayudaron a mejorar la legibilidad del artículo (en orden alfabético): Britney Muller, Douwe Kiela, Jared Casper, Jeff Rasley, Julien Launay, Leandro von Werra, Omar Sanseviero, Stefan Schweter y Thomas Wang.

Los gráficos principales fueron creados por Chunte Lee.

We will continue to update Zepes; if you have any questions or suggestions, please contact us!