Los investigadores de NVIDIA han desarrollado una nueva tecnología llamada compresión de memoria dinámica (DMC) que puede aumentar enormemente la eficiencia de la implementación de LLM y ampliar sus horizontes a secuencias más largas sin quedarse sin memoria.
A pesar del éxito de los modelos de lenguaje de gran tamaño (LLM) como herramientas de IA de propósito general, su alta demanda de recursos computacionales hace que su implementación sea un desafío en muchos escenarios del mundo real. Los tamaños del modelo y el estado de la conversación están limitados por la memoria de gran ancho de banda disponible, lo que limita la cantidad de usuarios que pueden ser atendidos y la duración máxima de la conversación.
En la actualidad, las dos arquitecturas LLM más populares son los transformadores y los modelos de espacio de estados selectivos (SSM), que se encuentran en los extremos opuestos de un espectro:
- Transformadores: El estado de conversación consiste en una representación distinta para cada elemento de una secuencia, cuyo tamaño aumenta rápidamente.
- SSM: comprime toda la secuencia en una única representación, que puede olvidar información anterior debido a su capacidad finita. La compresión del estado de conversación libera memoria y es esencial para ejecutar modelos más grandes dentro de las mismas limitaciones de memoria, procesar más tokens a la vez o simplemente reducir la latencia.
Para ello, los investigadores de NVIDIA han desarrollado una nueva tecnología llamada compresión de memoria dinámica (DMC) que puede aumentar enormemente la eficiencia de la implementación de LLM y ampliar sus horizontes a secuencias más largas sin quedarse sin memoria.
DMC abre una tercera vía, en la que se puede entrenar un modelo Transformer para comprimir de forma adaptativa el estado de la conversación y lograr una tasa de compresión deseada. Esto permite una reducción significativa del tamaño del estado de la conversación sin reemplazar la arquitectura Transformer habitual. DMC no requiere entrenamiento desde cero, ya que los modelos existentes se pueden adaptar mediante una cantidad insignificante de entrenamiento adicional, lo que es más confiable que los métodos de entrenamiento sin errores propensos a errores.
¿Qué afecta el rendimiento de la inferencia LLM?
La inferencia LLM consta de dos fases:
- Prellenado: se ingiere una consulta de usuario.
- Generación autorregresiva: la respuesta se genera un token a la vez.
Durante la generación, para realizar la autoatención, los transformadores agregan un par de representaciones (par clave-valor o KVP) para cada token a una memoria caché. Se almacena un KVP diferente para cada capa y cada cabezal de atención.
Como resultado, la caché KVP crece proporcionalmente a la longitud de la secuencia. Como la caché KVP debe caber en la memoria de la GPU junto con los pesos LLM, puede ocupar una parte significativa de ella o incluso agotarla.
Además, cuanto mayor sea el tamaño de la caché de KVP, más tiempo llevará ejecutar un único paso de inferencia. Esto se debe a que el cálculo de las puntuaciones de atención es una operación limitada a la memoria. Cada consulta tiene su propia caché de KVP que se debe cargar.
La situación es diferente para las proyecciones lineales en capas de atención o FFN, donde cada matriz de peso debe cargarse en SRAM desde HBM una vez para todas las consultas, si la GPU está trabajando en muchas consultas al mismo tiempo en paralelo.
En investigaciones anteriores se intentó reducir el tamaño de la caché de KVP cuantificando sus representaciones, compartiendo los puntos de atención o eliminando tokens de ella. Sin embargo, estos métodos degradan el rendimiento original porque eliminan información de la memoria sin alterar el comportamiento original de LLM.
Compresión de memoria dinámica
La compresión de memoria dinámica (DMC) es una forma sencilla de comprimir la caché KV durante la inferencia sin incurrir en una caída del rendimiento. El modelo decide por separado para cada token, capa y núcleo si el nuevo par de clave y valor se debe agregar a la caché KV, como es habitual en Transformer simple, o acumularse en el último, utilizando la siguiente ecuación:
Esta ecuación, que se encuentra en el corazón de DMC, transforma una subsecuencia de claves en una suma de prefijo particular, que recuerda a SSM populares como xLSTM o RWKV.
Durante la inferencia, los valores de alfa son estrictamente binarios. Cuando 
En un modelo simple, la caché se amplía con un KVP a la vez. Con DMC, una variable de decisión determina si se debe ampliar la caché o si el nuevo par se debe fusionar con el último en la caché de KVP.
Un componente crucial del DMC es el proceso de modernización , que enseña a los LLM este comportamiento:
- Entrene LLM preexistentes, como los de la familia Llama, utilizando entre el 2 y el 8 % de la mezcla de datos de entrenamiento original.
- Inicialice los módulos de decisión para que siempre agreguen KVP a la caché de KVP
- Realice una transición lenta hacia la DMC ejerciendo presión para promediar los pares nuevos con los pares posteriores. La tasa de compresión objetivo aumenta de 1x al nivel deseado durante el proceso de modernización.
- Después de alcanzar la tasa de compresión objetivo, fíjela para los pasos finales de modernización para consolidarla.
La decisión de añadir o fusionar es discreta. Para entrenar LLM con descenso de gradiente, se realiza una relajación continua de esta decisión a través de la distribución Gumbel-Sigmoid, lo que da como resultado elementos de memoria parcialmente añadidos y parcialmente fusionados durante el entrenamiento.
Para simular el DMC en tiempo de inferencia (Figura 2), donde los tokens no pueden asistir a etapas intermedias de fusión, se aumenta la máscara causal de los Transformadores para evitar cada vez más el acceso a esos estados.
En general, este procedimiento le permite evitar picos de perplejidad durante el preentrenamiento continuo. De hecho, descubrimos que enfoques como el de atención a consultas generales (GQA) , que modifican significativamente la arquitectura del modelo antes del preentrenamiento continuo, tienen dificultades para recuperar el rendimiento original.
Resultados
DMC logra un rendimiento comparable con los modelos vanilla para todas las aplicaciones posteriores que evaluamos, incluidas la factualidad ( MMLU ), la respuesta a preguntas con sentido común y la codificación ( HumanEval ).
Fundamentalmente, DMC se puede combinar con otros métodos para la reducción de caché KVP, como la cuantificación y GQA.
| Tasa de compresión DMC | MMLU (factibilidad) | Control de calidad basado en el sentido común | HumanEval (codificación) | |
| Llama-2-7B | 1x | 44.6 | 70.5 | 14.0 |
| Llama-2-7B | 4x | 44.2 | 70.2 | 16.5 |
| Llama-2-7B | 8x | 41.8 | 70.1 | 16.5 |
| Llama-2-13B | 1x | 54.5 | 73.5 | 17.5 |
| Llama-2-13B | 4x | 54.2 | 73.2 | 22.0 |
| Llama-2-13B | 8x | 52.1 | 73.3 | 21.3 |
| Llama-2-70B + AGC | 8x | 68.8 | 78.0 | 29.6 |
| Llama-2-70B + AGC | 16x | 68.8 | 77,9 | 29.9 |
Tabla 1. Precisión de la tarea posterior durante la inferencia de modelos vainilla (tasa de compresión 1x) y modelos habilitados para DMC
| Llama-2-7B + DMC 4x | Llama-2-7B + DMC 4x (cuantificación de 8 bits) | Llama-2-7B + DMC 8x | L lama-2-7B + DMC 8x (cuantificación de 8 bits) | |
| MMLU (factibilidad) | 44.2 | 44.6 | 41.8 | 41.7 |
Tabla 2. Puntuaciones MMLU de modelos habilitados para DMC con cuantificación FP8 aplicada a pesos, activaciones y caché KVP
En algunos casos, incluso se mejora el rendimiento. Hemos especulado que esto se debe a fenómenos que mitigan la compresión, como los tokens perdidos en el medio , aunque se necesitan más pruebas para corroborarlo.
La compresión del estado de conversación libera memoria que podría utilizarse para alojar un tamaño de lote mayor. Como el modelo está limitado por la memoria, esto suele traducirse casi directamente en un aumento equivalente del rendimiento.
Por ejemplo, elegir el tamaño de lote máximo que cabe en la memoria con una compresión de 8x en una GPU NVIDIA H100 significa un 700 % más de tokens generados por segundo que con el modelo estándar.
Después de una cierta tasa de compresión, el cálculo de las puntuaciones de atención queda limitado por el cálculo, por lo que la latencia puede aumentar mientras se atienden más consultas.
En la Figura 4, el eje x muestra el tamaño máximo de lote que cabe en la memoria en una sola GPU (7B y 13B) o dos GPU con paralelismo tensorial (70B) para el LLM estándar y los modelos DMC 2x, 4x, 6x y 8x.
Si el tamaño del lote permanece invariable, la latencia disminuye. En la práctica, ajustar el tamaño del lote permite optimizar el equilibrio entre latencia y rendimiento.
En la Figura 5, hasta un determinado tamaño de lote, los modelos están limitados por la memoria y pueden procesar ejemplos adicionales con la misma latencia. Más allá de este punto, los modelos se vuelven cada vez más limitados por el cálculo. Las mediciones se tomaron durante la generación de los últimos 1K tokens de una secuencia de 4K tokens en una GPU NVIDIA H100.
Debido a la flexibilidad de comprimir diferentes cabezas y capas de atención a diferentes velocidades, DMC también revela nuevos conocimientos sobre el funcionamiento interno de los LLM. Por ejemplo, las capas superiores son las que más se comprimen, ya que sus KVP son los más redundantes.
En la Figura 6, las cabezas están dispuestas desde la compresión más alta hasta la más baja, de arriba hacia abajo para mayor claridad.
Conclusión
DMC promete llevar las fronteras de los LLM aún más lejos.
Proporciona un mecanismo para la memoria adaptativa, que se encuentra entre la memoria lineal en los Transformers y la memoria constante de los SSM, y ofrece un mejor equilibrio entre capacidad y espacio. Esto aumenta el rendimiento de Transformer LLM sin sacrificar la calidad y permite la adaptación de contextos mucho más largos dentro de las mismas limitaciones de hardware. Esto es crucial para el escalamiento del tiempo de inferencia a través de la cadena de pensamiento, que requiere generar largas secuencias de tokens.
DMC también abre nuevas vías para crear de forma económica prototipos y LLM a escala de producción con nuevas arquitecturas, reutilizando las preexistentes.
Para obtener más información, consulte los siguientes recursos:
Expresiones de gratitud
Nos gustaría agradecer a Mostofa Patwary y Szymon Migacz por su ayuda, así como a Przemysław Strzelczyk, Daniel Korzekwa y Bryan Catanzaro por las útiles discusiones y el apoyo para la publicación de este artículo. Este trabajo fue financiado en parte por el Centro de Formación Doctoral en Procesamiento del Lenguaje Natural del UKRI, financiado por el UKRI (subvención EP/S022481/1) y la Facultad de Informática y la Facultad de Filosofía, Psicología y Ciencias del Lenguaje de la Universidad de Edimburgo.
Recursos relacionados
- Sesión de GTC: Aceleración de aplicaciones de aprendizaje profundo con compresión sobre la marcha basada en GPU
- Sesión de GTC: Más datos, más rápido: prácticas recomendadas de gestión de memoria de GPU en Python y C++
- Sesión de GTC: Estrategias avanzadas de optimización del almacenamiento para cargas de trabajo de IA de producción dinámica (presentada por DDN)
- SDK: VDB neuronal
- Kit de desarrollo de software: nvCOMP
- SDK: complemento de Photoshop para herramientas de textura de NVIDIA
NVIDIA Blog. A. K., J. R. y D. T. Traducido al español
Contacto