Adaptando el diseño del sistema para grandes modelos de lenguaje

Alinear el diseño de sistemas futuros con las crecientes necesidades computacionales de los grandes modelos de lenguaje (LLMs) es un reto importante hoy en día. Se propone un método general para el modelado de rendimiento y análisis de carga de trabajo en el entrenamiento y la inferencia distribuidos de LLM. Este método toma en cuenta varios factores, incluyendo recursos computacionales, sistemas de memoria, capacidades de red y diferentes formas de paralelizar tareas. Estos métodos de paralelización incluyen Paralelismo de modelo, paralelismo de datos, paralelismo de tuberías y paralelismo de secuencias.

Para apoyar nuestras predicciones, validamos nuestros hallazgos usando datos de estudios de investigación y benchmarks de la industria. Para el entrenamiento distribuido, nos centramos en el uso de memoria de los LLM con diferentes métodos de recomputación de activaciones. Investigamos por qué hay un aumento notable en el rendimiento al pasar de GPUs A100 a B200 (un aumento de 35x), lo que sigue de cerca las tendencias proporcionadas por NVIDIA. También realizamos un estudio en diferentes nodos tecnológicos, que van desde 12 nm a 1 nm, para ver cómo las mejoras en lógica, memoria y red afectan el rendimiento. En términos de inferencia, analizamos el equilibrio entre recursos computacionales y uso de memoria para diferentes operaciones en varios sistemas de GPU y consideramos cómo los avances en la tecnología de memoria DRAM influyen en la velocidad de inferencia. Al usar nuestro enfoque de modelado de rendimiento, podemos identificar dónde surgen problemas de rendimiento tanto en el entrenamiento como en la inferencia de LLM a medida que la tecnología evoluciona. Esta información es crítica para diseñar sistemas futuros centrados en el entrenamiento y la inferencia de LLM.

#Introducción y Antecedentes

La arquitectura transformadora se ha convertido en un diseño estándar para diversas aplicaciones en inteligencia artificial (IA). Los grandes modelos de lenguaje (LLMs), como los modelos GPT y sus variantes, ejemplifican la escala y la complejidad de las arquitecturas transformadoras. Entrenar estos modelos requiere enormes cantidades de datos y recursos computacionales, lo que conlleva costos significativos y emisiones de carbono. Por ejemplo, entrenar un modelo como GPT-3 puede costar alrededor de 10 millones de dólares. Sin embargo, una parte considerable del costo total puede surgir de la inferencia cuando estos modelos sirven a muchos usuarios durante un largo periodo.

Para entender las tendencias en los LLM y su hardware subyacente, es crucial evaluar su rendimiento en relación con el costo total de operación (TCO) para el entrenamiento y la inferencia a gran escala. Un examen exhaustivo del rendimiento en relación con el TCO puede ayudar a señalar desafíos y priorizar inversiones en sistemas y modelos computacionales futuros.

#La Arquitectura Transformadora

La arquitectura transformadora basada en decodificadores consta de capas que principalmente presentan un bloque de atención multicanal (MHA) y un bloque de perceptrón multicapa (MLP). Esta estructura consistente permite analizar el rendimiento de los LLM a nivel de un gran centro de datos. El mecanismo de atención es central en el diseño transformador, donde el éxito de los LLM depende de su longitud de secuencia o contexto. Sin embargo, aumentar la longitud de la secuencia resulta en mayores demandas de ejecución y memoria.

Tres factores clave afectan el rendimiento: el número de operaciones de punto flotante (FLOPs), accesos a memoria y comunicación de datos. Iniciativas recientes, como FlashAttention, abordan problemas de acceso a la memoria optimizando cómo se recuperan los datos de la DRAM, potencialmente a expensas de los FLOPs. Además, el uso de KV-cache es vital para mejorar el rendimiento de inferencia. Los desafíos de rendimiento no son estáticos; cambian a medida que los LLM y el hardware computacional avanzan. Estudios indican que la comunicación puede representar un porcentaje significativo (40-75%) del tiempo de ejecución a medida que los modelos y el hardware evolucionan. Por lo tanto, es necesario un marco amplio que resalte las compensaciones en el rendimiento durante el entrenamiento y la inferencia de LLM para una colaboración efectiva entre hardware y software.

#Cuellos de Botella de Rendimiento

Las operaciones transformadoras se clasifican en tres categorías principales: contracciones tensoriales (como GEMM o GEMV), normalización (como softmax o normalización de capas) y operaciones elemento a elemento (incluyendo funciones como no linealidades o sesgos). La operación más crítica suele ser GEMM o GEMV, que puede impactar significativamente el rendimiento general de un modelo transformador.

El rendimiento de GEMM puede analizarse comparando el tiempo tomado para el cálculo frente a los accesos a memoria. Si el tiempo de cálculo domina, se considera que la operación está limitada por el cómputo, mientras que si el tiempo de acceso a memoria es mayor, la operación se clasifica como limitada por memoria. La intensidad aritmética es una medida que indica si una tarea está limitada por el cómputo o por la memoria según el número de operaciones aritméticas relacionadas con las transferencias de memoria.

En el entrenamiento distribuido, las contracciones tensoriales suelen ser intensivas en cómputo debido a los GEMMs "gruesos". En contraste, durante la fase de inferencia, las operaciones tienden a estar limitadas por memoria debido a la naturaleza menos intensiva de la generación de tokens, resultando en GEMMs "delgados". La fusión de kernels es una técnica común para mejorar la intensidad aritmética de las operaciones limitadas por memoria.

El segundo aspecto esencial para el entrenamiento o la inferencia de LLM es la transferencia de datos a través de la red. A medida que el entrenamiento y la inferencia escalan, los costos de comunicación se vuelven críticos, particularmente en sistemas con múltiples GPUs.

#Estrategias de Paralelización

El entrenamiento distribuido aprovecha varias estrategias de paralelización como el paralelismo de datos (DP), el paralelismo tensorial de modelos (TP), el paralelismo de tuberías (PP) y el paralelismo de secuencias (SP). En DP, cada GPU procesa una porción de los datos mientras comparte los mismos parámetros del modelo para calcular gradientes locales. Estos gradientes se agregan para actualizar los parámetros del modelo.

Los requerimientos de memoria en DP dependen de factores como el tamaño del lote, la longitud de la secuencia y las dimensiones del modelo. TP ayuda a reducir los requisitos de memoria relacionados con los parámetros del modelo al distribuir las operaciones tensoriales en múltiples dispositivos, permitiendo que cada dispositivo procese solo una parte de una matriz.

El paralelismo de tuberías divide las capas entre dispositivos, donde cada uno procesa un conjunto de capas y pasa los resultados. Este enfoque puede llevar a tiempos de inactividad, también conocidos como burbujas de tubería, que pueden mitigarse con técnicas de programación avanzadas. El paralelismo de secuencias emplea un enfoque similar para paralelizar ciertas operaciones, reduciendo su huella de memoria sin incurrir en retrasos de comunicación.

Entrenar LLM a menudo requiere combinaciones de estas estrategias de paralelización, mientras que la inferencia típicamente utiliza TP en menos dispositivos debido al tamaño de datos limitado.

#Marco para el Modelado de Rendimiento

Nuestro marco construye un enfoque general para modelar y analizar el rendimiento en cargas de trabajo de entrenamiento e inferencia de LLM distribuidos. Considera factores como sistemas de cómputo y memoria, así como la comunicación en red.

Después de una validación exhaustiva contra conjuntos de datos relacionados con GEMM, GEMV, entrenamiento e inferencia, podemos brindar información sobre las ganancias de rendimiento logradas al cambiar de GPUs A100 a B200. También exploramos opciones de diseño en diferentes nodos tecnológicos, observando cómo los avances en tecnologías DRAM impactan el rendimiento.

#Análisis de la Huella de Memoria

Al entrenar grandes LLM, una preocupación significativa es el desbordamiento de memoria. Comprender las características del uso de memoria es crucial para planear estrategias de entrenamiento efectivas. Nuestro perfilado de memoria examina tres métodos diferentes de recomputación de activaciones para optimizar la eficiencia del entrenamiento: sin recomputación, recomputación selectiva y recomputación completa.

Al perfilar la memoria para varios modelos GPT con estos métodos de recomputación, observamos las implicaciones en el uso de memoria. Sin recomputación, ajustar un LLM dentro de la memoria de un dispositivo a menudo resulta imposible a menos que usemos tamaños de lote extremadamente pequeños o un alto grado de paralelismo. Las opciones de recomputación selectivas pueden proporcionar una eficiencia de memoria ligeramente mejor con un mínimo de sobrecarga computacional.

#Escalado del Rendimiento a Través de Generaciones de GPU

La necesidad de procesamiento acelerado en cargas de trabajo de IA ha llevado a avances significativos en la tecnología de GPUs de NVIDIA. Usando nuestro modelo de rendimiento, podemos proyectar cómo el tiempo de entrenamiento para un modelo sustancial, como GPT-3, varía a través de diferentes generaciones de GPUs.

Hay una clara tendencia que demuestra el escalado del rendimiento de A100 a H100 y de H100 a B200, con mejoras en gran parte vinculadas a avances en capacidad de DRAM y potencia de procesamiento. Cada nueva generación ha introducido características que mejoran la capacidad de cómputo y el ancho de banda de memoria, lo que lleva a disminuciones notables en el tiempo de entrenamiento.

#Escalado de la Tecnología de Memoria

A medida que la tecnología de memoria mejora, analizamos cómo estos avances afectan la velocidad de inferencia, particularmente en sistemas multi-GPU. Al estudiar cómo diferentes tecnologías DRAM impactan el rendimiento, podemos hacer observaciones sobre cómo el ancho de banda de memoria se correlaciona con el tiempo de inferencia.

Nuestros hallazgos indican que, hasta cierto punto, el rendimiento mejora linealmente con el aumento del ancho de banda de DRAM. Sin embargo, a medida que las tecnologías avanzan, este escalado lineal puede disminuir, y otros factores como la capacidad y el ancho de banda de la memoria en chip se convierten en componentes críticos que influyen en el rendimiento general.

#Conclusión

En resumen, hemos examinado las cargas de trabajo de entrenamiento e inferencia de LLM a través de un marco detallado de modelado de rendimiento. El enfoque nos permite analizar varios casos de prueba y evaluar el impacto del escalado de la tecnología computacional y de memoria en el rendimiento.

Al evaluar estos parámetros, logramos predicciones precisas tanto para el entrenamiento como para la inferencia y podemos identificar los cuellos de botella clave de rendimiento que surgen con los avances en tecnología y arquitectura. Los esfuerzos futuros incluyen refinar nuestro enfoque para comprender mejor la utilización de memoria en escenarios más intensivos en memoria e incorporar simulaciones de red para explorar el impacto del tráfico de datos en el rendimiento general. Además, planeamos evaluar las próximas arquitecturas y tecnologías de sistemas para el entrenamiento e inferencia de LLM, centrándonos en la eficiencia energética y el análisis del costo total de operación.

Esta investigación tiene como objetivo proporcionar valiosos insights para optimizar el diseño y la implementación de sistemas de entrenamiento e inferencia de LLM, asegurando que satisfagan de manera efectiva las crecientes demandas computacionales.