Introducción a las optimizaciones de inferencia de modelos de lenguaje grandes (LLM): 1. Antecedentes y formulación del problema

Descripción general de la inferencia del modelo de lenguaje grande (LLM), su importancia, desafíos y formulaciones de problemas clave.

Los modelos de lenguaje de gran tamaño (LLM, por sus siglas en inglés) han revolucionado el campo del procesamiento del lenguaje natural (NLP, por sus siglas en inglés) al permitir una amplia gama de aplicaciones, desde chatbots y agentes de IA hasta generación de código y contenido. Sin embargo, la implementación de LLM en escenarios del mundo real a menudo enfrenta desafíos relacionados con la latencia, el consumo de recursos y la escalabilidad.

En esta serie de publicaciones del blog, exploraremos varias técnicas de optimización para la inferencia LLM. Analizaremos en profundidad estrategias para reducir la latencia, el uso de memoria y el costo computacional, desde mecanismos de almacenamiento en caché hasta aceleraciones de hardware y cuantificación de modelos.

En esta publicación, brindaremos una breve descripción general de la inferencia LLM, su importancia y los desafíos asociados con ella. También describiremos las formulaciones de problemas clave que guiarán nuestra exploración de las técnicas de optimización.

Inferencia de modelos: una descripción general

La inferencia de modelos se refiere al proceso de utilizar un modelo de aprendizaje automático entrenado para hacer predicciones o generar resultados basados en nuevos datos de entrada. En el contexto de los LLM, la inferencia implica procesar la entrada de texto y generar una salida de texto coherente y contextualmente relevante.

El modelo se entrena solo una vez o periódicamente, mientras que la inferencia ocurre con mucha más frecuencia, probablemente miles de veces por segundo en entornos de producción.

La optimización de la inferencia es esencial para garantizar que los LLM se puedan implementar de manera eficaz en aplicaciones del mundo real. El objetivo es minimizar la latencia (el tiempo que se tarda en generar una respuesta), reducir el consumo de recursos (CPU, GPU, memoria) y mejorar la escalabilidad (la capacidad de manejar cargas crecientes).

Por ejemplo, GPT-3 (con 175 mil millones de parámetros) requiere recursos computacionales significativos para la inferencia. Las optimizaciones pueden reducir los tiempos de respuesta de 1 a 2 segundos a milisegundos, lo que hace que los LLM sean más prácticos para aplicaciones interactivas.

Descripción general de la arquitectura del transformador

La arquitectura de transformadores, que utiliza mecanismos de atención, se ha convertido en la base de la mayoría de los LLM de última generación. Esta arquitectura incluye codificaciones posicionales, autoatención de múltiples cabezas, redes neuronales de avance y normalización de capas. Los transformadores se clasifican generalmente en tres tipos principales:

• Los modelos que solo utilizan codificadores (por ejemplo, BERT) están diseñados para tareas como la clasificación de texto y el reconocimiento de entidades con nombre. Convierten una secuencia de entrada en una representación de longitud fija (incrustación). Estos modelos son bidireccionales, lo que significa que tienen en cuenta el contexto tanto desde la izquierda como desde la derecha de un token, lo que puede llevar a una mejor comprensión del texto de entrada.

• Los modelos que solo utilizan decodificadores (por ejemplo, GPT-3) se utilizan para tareas de generación de texto. A partir de una secuencia de entrada, generan texto de un token a la vez, condicionándolos a los tokens generados previamente. Estos modelos son unidireccionales, lo que significa que solo consideran el contexto desde la izquierda de un token, lo que resulta adecuado para tareas como el modelado de lenguaje. Esta es la arquitectura LLM más común.

• Los modelos de codificador-decodificador (por ejemplo, T5) fueron la arquitectura original que se presentó en el artículo “Attention is All You Need”. Estos modelos están diseñados para tareas que requieren tanto comprensión como generación, como traducción y resumen. Procesan la secuencia de entrada con el codificador y luego generan la secuencia de salida con el decodificador.

  
  

Dado que los modelos de solo decodificador son la arquitectura LLM más común para tareas autorregresivas, esta serie se centrará en técnicas de optimización específicamente para este tipo de modelo.

Descripción general del mecanismo de atención

El mecanismo de atención es un componente clave de la arquitectura del transformador que permite que el modelo se centre en diferentes partes de la secuencia de entrada al generar la salida. Calcula una suma ponderada de las representaciones de entrada, donde los pesos se determinan según la relevancia de cada token de entrada con respecto al token de salida actual que se está generando. Este mecanismo permite que el modelo capture dependencias entre tokens, independientemente de su distancia en la secuencia de entrada.

El mecanismo de atención puede ser costoso en términos computacionales, especialmente para secuencias de entrada largas, ya que requiere calcular interacciones por pares entre todos los tokens (complejidad O(n^2) ). Veámoslo con más detalle paso a paso:

1. Representación de entrada : cada token en la secuencia de entrada se representa como un vector, generalmente mediante incrustaciones.

   
   

2. Vectores de consulta, clave y valor : para cada token, se calculan tres vectores: un vector de consulta ( Q_i ), un vector de clave ( K_i ) y un vector de valor ( V_i ). Estos vectores se derivan de las representaciones de entrada mediante transformaciones lineales aprendidas.

   
   

3. Puntuaciones de atención : las puntuaciones de atención se calculan tomando el producto escalar del vector de consulta del token actual con los vectores clave de todos los tokens anteriores en la secuencia de entrada. Esto da como resultado una puntuación que indica cuánta atención se debe poner en cada token.

   
   

4. Normalización Softmax : luego, los puntajes de atención se normalizan utilizando la función softmax para obtener pesos de atención, que suman 1.

   
   

5. Suma ponderada : finalmente, la representación de salida del token actual se calcula como una suma ponderada de los vectores de valor, utilizando los pesos de atención.

Atención de múltiples cabezas

La atención multicabezal es una extensión del mecanismo de atención que permite que el modelo preste atención de manera conjunta a la información de diferentes subespacios de representación en diferentes posiciones. En lugar de tener un único conjunto de ponderaciones de atención, la atención multicabezal calcula varios conjuntos de puntuaciones de atención en paralelo, cada uno con sus propias transformaciones lineales aprendidas.

Las salidas de estos cabezales de atención se concatenan y se transforman linealmente para producir la representación de salida final.

Este mecanismo mejora la capacidad del modelo para capturar diversas relaciones y dependencias en los datos de entrada, lo que conduce a un mejor rendimiento en varias tareas de PNL.

Descripción general del proceso de cálculo de inferencia

Con una comprensión de los LLM y la arquitectura de los transformadores, describamos el proceso de cálculo de inferencia. La inferencia genera los siguientes $n$ tokens para una secuencia de entrada determinada y se puede dividir en dos etapas:

1. Etapa de prellenado : en esta etapa, se realiza un paso hacia adelante a través del modelo para la secuencia de entrada y se calculan las representaciones de clave y valor para cada token. Estas representaciones se almacenan para su uso posterior en la etapa de decodificación en una caché KV. Las representaciones de todos los tokens en cada capa se calculan en paralelo.

   

2. Etapa de decodificación : en esta etapa, el modelo genera los tokens de salida de a uno por vez de manera autorregresiva. Para cada token, el modelo obtiene las representaciones de clave y valor de la memoria caché KV almacenada durante la etapa de prellenado, junto con la representación de consulta del token de entrada actual para calcular el siguiente token en la secuencia.

   
   

   Este proceso continúa hasta que se cumple un criterio de detención (por ejemplo, alcanzar una longitud máxima o generar un token de fin de secuencia). Las nuevas representaciones de clave y valor se almacenan en la caché KV para los tokens subsiguientes. En esta etapa, también se aplica una estrategia de muestreo de tokens para determinar el siguiente token que se generará (por ejemplo, búsqueda voraz, búsqueda de haz, muestreo top-k).

   

Complejidad del cálculo de inferencia

Para un prefijo de longitud L , tamaño de incrustación d y un modelo con h cabezas y n capas, la complejidad del cálculo de inferencia se puede analizar de la siguiente manera:

• Etapa de prellenado : en la etapa de prellenado, calculamos la representación inicial de todos los tokens en la entrada. La complejidad aquí es:

  Aquí:

  • Primer término O(Ln .d^2) : representa el cálculo de avance, que procesa cada token de forma independiente en las distintas capas. Esto se escala linealmente con la longitud de la secuencia L y la cantidad de capas n .

  
  

  • Segundo término O(L^2. nh d) : representa el costo del mecanismo de atención. Aquí, cada token interactúa con todos los demás tokens, lo que da como resultado una complejidad L^2 para el cálculo de la atención por capa. La complejidad crece cuadráticamente con la longitud de la secuencia, lo que puede convertirse en un importante cuello de botella para secuencias largas.

    
    

• Etapa de decodificación : La etapa de decodificación es la parte autorregresiva, la complejidad es:

• Aquí:

  • Cálculo de avance : para cada token generado, realizamos operaciones de avance en cada capa. Dado que se realiza para un token a la vez (no para toda la secuencia), la complejidad por token es: O(nd^2) .

    
    

  • Cálculo de la atención con almacenamiento en caché : cada nuevo token interactúa con la secuencia existente a través de la atención, utilizando los pares clave-valor calculados previamente. Para cada token generado, este cálculo de la atención es proporcional a la longitud de la secuencia L, lo que da como resultado: O(Lnd .h)

    
    

Como podemos ver, la complejidad del cálculo de inferencia está influenciada por la longitud de la secuencia de entrada ( L ), el número de capas ( n ), el número de cabezas de atención ( h ) y el tamaño de incrustación ( d ). Esta complejidad puede convertirse en un cuello de botella en aplicaciones en tiempo real, especialmente cuando se trabaja con secuencias de entrada largas y/o modelos grandes.

Importancia del almacenamiento en caché de KV

El almacenamiento en caché de KV es una técnica de optimización crucial para la inferencia LLM, en particular en la etapa de decodificación. Al almacenar las representaciones de clave y valor calculadas durante la etapa de prellenado, el modelo puede evitar cálculos redundantes para tokens procesados previamente.

Esto reduce significativamente el costo computacional y la latencia durante la inferencia, ya que el modelo solo necesita calcular los puntajes de atención para el nuevo token que se genera, en lugar de volver a calcular las representaciones de clave y valor para todos los tokens en la secuencia de entrada.

Esto hace que el costo sea lineal con respecto a la cantidad de tokens generados, en lugar de cuadrático con respecto a la longitud de entrada.

Sin embargo, el almacenamiento en caché de KV requiere memoria adicional para almacenar las representaciones de claves y valores, lo que puede ser una desventaja en entornos con recursos limitados.

Cálculos para un modelo de ejemplo

Calculemos los requisitos de memoria para el modelo LLaMA 7B.

Configuración del modelo

• Parámetros: 7 mil millones
• Tamaño de incrustación ( d_model ): 4096
• Número de capas: 32
• Número de cabezas de atención ( d_head ): 32
• Dimensión de la cabeza ( d_head ): 128 (4096/32)
• Longitud máxima de secuencia (L): 2048
• Tipo de dato: float16 (2 bytes por elemento)

Cálculo de memoria

1. Tamaño de caché por capa : para cada capa, necesitamos almacenar tanto claves como valores

   • Tamaño de clave por token = d_head × num_heads = 128 × 32 = 4096 elementos

   • Tamaño del valor por token = d_head × num_heads = 128 × 32 = 4096 elementos

   • Total de elementos por token por capa = 4096 + 4096 = 8192 elementos

     
     

2. Memoria por capa para secuencia completa : para la secuencia completa de longitud L = 2048 tokens

   • Elementos por capa = L × 8192 = 2048 × 8192 = 16.777.216 elementos

   • Memoria por capa (en bytes) = 16.777.216 × 2 = 33.554.432 bytes = 33,55 MB

     
     

3. Memoria caché KV total para todas las capas : dado que tenemos 32 capas
   • Memoria total = 33,55 × 32 MB = 1073,6 MB

Requerimiento total de memoria

• Pesos del modelo: 7 mil millones de parámetros × 2 bytes/parámetro = 14 GB

• Memoria caché KV: 1073,6 MB

• Otros recursos de memoria (por ejemplo, activaciones, resultados intermedios): ~1-2 GB

Por lo tanto, el requerimiento total de memoria es: 14 GB (pesos del modelo) + 1-2 GB (sobrecarga) + 1073,6 MB (caché KV) = 15-16 GB . Este cálculo nos da una estimación de los requerimientos de memoria para el modelo LLaMA 7B durante la inferencia. LLaMA 7B es relativamente pequeño en comparación con modelos como GPT-3 (175 mil millones de parámetros), que requerirían significativamente más memoria tanto para los pesos del modelo como para el caché KV.

Además, si se escala a $m$ usuarios simultáneos, los requisitos de recursos serían $m$ veces mayores. Por lo tanto, las técnicas de optimización son cruciales para implementar modelos grandes en entornos con recursos limitados.

Métricas para evaluar la optimización de la inferencia

Al evaluar la eficacia de las técnicas de optimización de inferencia, se pueden considerar varias métricas:

1. Latencia de pre-relleno : el tiempo que se tarda en realizar la etapa de pre-relleno de la inferencia, también llamada latencia de tiempo hasta el primer token (TTFT). Esta métrica es crucial para las aplicaciones interactivas donde los usuarios esperan respuestas rápidas. Factores como el tamaño del modelo, la longitud de entrada y las capacidades del hardware pueden influir en esta métrica.

   
   

2. Latencia de decodificación : el tiempo que se tarda en generar cada token posterior a la etapa de prellenado, también llamada latencia entre tokens (ITL). Esta métrica es importante para medir la capacidad de respuesta del modelo durante la generación de texto. Para aplicaciones como los chatbots, una ITL baja es buena, pero una velocidad más rápida no siempre es mejor, ya que entre 6 y 8 tokens por segundo suelen ser suficientes para la interacción humana. Los factores que influyen incluyen el tamaño de la caché KV, la estrategia de muestreo y el hardware.

   
   

3. Latencia de extremo a extremo : el tiempo total que transcurre desde la recepción de la entrada hasta la generación del resultado final. Esta métrica es esencial para comprender el rendimiento general del proceso de inferencia y se ve influenciada por el prellenado, la decodificación y otras latencias de los componentes (por ejemplo, el análisis de JSON). Los factores que influyen incluyen el tamaño del modelo, la longitud de la entrada y el hardware, así como la eficiencia de todo el proceso.

   
   

4. Tasa máxima de solicitudes, también conocida como QPS (consultas por segundo) : la cantidad de solicitudes de inferencia que se pueden procesar por segundo. Esta métrica es crucial para evaluar la escalabilidad del modelo en entornos de producción. Factores como el tamaño del modelo, el hardware y las técnicas de optimización pueden influir en la QPS. Por ejemplo, si se ofrecen 15 QPS para una latencia P90 a través de 1 GPU, para ofrecer 300 QPS, se necesitarían 20 GPU. Los factores que influyen incluyen los recursos de hardware, el equilibrio de carga y las técnicas de optimización.

   
   

5. FLOPS (operaciones de punto flotante por segundo) : la cantidad de operaciones de punto flotante que el modelo puede realizar en un segundo. Esta métrica es útil para comprender el costo computacional de la inferencia y se puede utilizar para comparar la eficiencia de diferentes modelos y técnicas de optimización. Los factores que influyen incluyen la arquitectura del modelo, el hardware y las técnicas de optimización.

Tipos de técnicas de optimización de inferencia

Cubriremos todas estas optimizaciones en futuras publicaciones de la serie.

• Optimización de la arquitectura del modelo : modificación de la arquitectura del modelo para mejorar la eficiencia de la inferencia, como reducir la cantidad de capas o cabezales de atención, o utilizar mecanismos de atención más eficientes (por ejemplo, atención dispersa).

• Optimización del sistema : optimización de la infraestructura de hardware y software subyacente, como el uso de hardware especializado (por ejemplo, TPU, GPU) u optimización de la pila de software (por ejemplo, el uso de bibliotecas y marcos eficientes). Se puede dividir en:

  
  

  • Gestión de memoria : gestionar de forma eficiente el uso de la memoria para reducir la sobrecarga y mejorar el rendimiento.

  • Computación eficiente : aprovechar el paralelismo y optimizar el cálculo para reducir la latencia.

  • Procesamiento por lotes : procesamiento de múltiples solicitudes simultáneamente para mejorar el rendimiento.

  • Programación : Programar tareas de manera eficiente para maximizar la utilización de recursos.

    
    

• Compresiones de modelos : se pueden utilizar técnicas como cuantificación, poda y destilación para reducir el tamaño del modelo y mejorar la velocidad de inferencia sin sacrificar significativamente el rendimiento.

• Optimización de algoritmos : mejora de los algoritmos utilizados para la inferencia, como el uso de estrategias de muestreo más eficientes o la optimización del mecanismo de atención. Por ejemplo, la decodificación especulativa, que permite que el modelo genere múltiples tokens en paralelo, puede reducir significativamente la latencia de decodificación.

  

  
  

Conclusión

En esta publicación, brindamos una descripción general de la inferencia LLM, su importancia y los desafíos asociados con ella. También describimos las formulaciones de problemas clave que guiarán nuestra exploración de las técnicas de optimización en publicaciones posteriores.

Al comprender las complejidades de la inferencia LLM y los factores que influyen en su rendimiento, podemos apreciar mejor la importancia de las técnicas de optimización para hacer que los LLM sean más prácticos para las aplicaciones del mundo real. En la próxima publicación, profundizaremos en técnicas de optimización específicas y sus implementaciones, centrándonos en reducir la latencia y el consumo de recursos mientras se mantiene el rendimiento del modelo.

Referencias

• Atención es todo lo que necesitas.
• Optimización de inferencias de modelos básicos en aceleradores de IA