Su caché KV no tiene un problema de bits, tiene un problema de geometría.

A una precisión de 2 bits idéntica, una decisión sobre qué eje se cuantiza hace que la puntuación de una prueba de referencia pase de 2,88 a 63,53. Las claves y los valores necesitan un tratamiento opuesto, y la razón está en la ecuación de atención, no en el hardware.

Tomemos el modelo Llama-2-13B. Agrupemos su caché de clave-valor por un tamaño de grupo de cuantización de 32 en dos bits, mientras que dejamos todo lo demás en su lugar, mismo modelo, mismo presupuesto de bits, mismo tamaño de grupo, mismas pruebas de referencia.

Dependiendo de un solo elemento de una decisión de implementación, los resultados de precisión de CoQA son de 2,88 o 63,53. La puntuación utilizando precisión completa es de 66,37.

La decisión no se trata de cuántos bits totales se utilizan. La pregunta es simplemente qué eje se elige para agrupar a través de cuando se calcula cada factor de escala. Cuando se decide utilizar el canal como dimensión de agrupación (claves) y el token como dimensión de agrupación (valores), se termina en algún lugar dentro de cuatro puntos de la precisión de precisión completa. Si se invierte cualquiera de esas opciones, se experimenta una pérdida de calidad. Si se invierten ambas opciones, el modelo ya no funciona.

Cuatro formas de gastar los mismos 2 bits en la misma caché. Resultados de la ablación KIVI en Llama-2-13B con un tamaño de grupo de 32.

La cuantización suele considerarse como un solo control: 8 bits, 4 bits, 2 bits, con un costo de precisión suave. Dentro de la caché KV, no es así. Se trata de elegir sistemas de coordenadas, y diferentes sistemas se aplican a claves y valores. Este artículo explica por qué. En resumen: el error de cuantización depende del rango de valores dentro de los grupos; las claves y los valores tienen una estructura muy diferente; y la gente a menudo se equivoca porque no se puede derivar el eje correcto de la distribución de valores en absoluto. Hay que mirar cómo cambia el error después de que la atención lo consume. Eso da un principio general para comprimir activaciones intermedias y una buena razón para dudar del error de reconstrucción como proxy de calidad.

Por qué la caché KV es donde esto muerde

Durante la fase de generación, un transformador almacena todos los datos de proyección de clave-valor (KV) de los tokens que ha procesado anteriormente en una caché para que no tenga que recalcular estos datos de nuevo. Esa caché crece linealmente con la longitud del contexto y el tamaño del lote. Eventualmente, esto hará que la caché crezca más grande que el modelo en sí.

Este aumento en el crecimiento se puede identificar fácilmente cuando se mira el consumo de memoria de las diferentes partes del modelo. En el análisis de KVQuant de LLaMA-7B, los pesos cuentan con aproximadamente el 98 por ciento de la memoria a una longitud de secuencia de 512, con activaciones en el 2 por ciento. A 128K de contexto, la relación se invierte a alrededor del 16 por ciento de pesos y el 84 por ciento de caché KV. Cuando miramos un análisis de OPT-175B citado por los autores de KIVI, encontraron resultados similares. En particular, a un tamaño de lote de 512 con un prompt de 512 tokens, la caché KV alcanza 1,2 TB, varias veces el tamaño de los pesos del modelo.

Sin embargo, la capacidad es solo la mitad del problema aquí. La GPU debe leer toda la caché KV desde la memoria del dispositivo para cada token que genera. Esto significa que mientras la GPU está leyendo la caché KV, los núcleos de cómputo están inactivos. Al reducir el tamaño general de la caché, se aumenta el procesamiento disponible y se reduce el tiempo que se pasa esperando a las transferencias de datos.

Qué es realmente el error de cuantización

La cuantización de enteros uniforme es matemáticamente sencilla. Para un grupo de números, se registra el número más pequeño como un punto cero y luego se divide el rango de ese grupo por la cantidad de niveles que se pueden representar para obtener un tamaño de paso. Luego, se redondea cada elemento al paso más cercano. Dos resultados inmediatos siguen. Primero, el error por elemento está limitado por la mitad de un paso. Segundo, el tamaño del paso es el rango del grupo dividido por 2ᴮ − 1. A 2 bits, solo hay 4 niveles para cubrir cualquier dispersión que exista dentro de ese grupo. Entonces, un elemento que es cien veces más grande en comparación con sus vecinos no solo se desempeña mal. Infla el tamaño del paso para todos los demás elementos que comparten el mismo grupo, y todos se vuelven más gruesos juntos. El grupo es la unidad de daño. Elegir un eje significa decidir qué elementos sufren juntos. Enmarcar la pregunta de manera diferente, ya no es “¿Cuántos bits puedo permitirme sacrificar?” sino “¿Dónde están los valores extremos y puedo aislarlos?”

Claves: los valores atípicos viven en canales fijos

Los grandes modelos de lenguaje contienen activaciones que son inusualmente grandes en comparación con la mayoría de las activaciones. Sun y sus colegas catalogaron estas activaciones muy grandes en diferentes familias de modelos: en Mixtral 8x7B, la magnitud más grande está cerca de 7000, mientras que la magnitud de la característica mediana es de alrededor de 0,3, aproximadamente cuatro órdenes de magnitud aparte. Estos son muy raros; permanecen fijos en dimensiones que rara vez cambian con la entrada, y no son accidentales. Actúan como sesgos implícitos, y son lo que enfoca la atención en solo unos pocos tokens: el comportamiento del sumidero de atención. En la caché de claves, esta estructura es muy clara: canales específicos llevan magnitudes muy grandes consistentemente en todos los tokens de una secuencia. Agrupar a lo largo de los tokens, y cada grupo contiene esos canales atípicos, por lo que el tamaño del paso de cada grupo se establece por los valores atípicos, y todos los canales normales pagan por ello. Agrupar a lo largo de los canales, y los canales atípicos forman sus propios grupos. Su rango interno es grande pero autocontenido; los canales normales se dejan solos. Los resultados coinciden. Promediados en capas y cabezas en Llama-2-13B, KIVI informa un error de reconstrucción de claves de 13,67 bajo agrupación por token frente a 4,55 por canal, y, lo que es más importante, un error de puntuación de atención de 47,00 frente a 9,60. La cuantización de claves por token produce aproximadamente cinco veces el error de puntuación. Las puntuaciones entonces coinciden con las métricas significativas para las claves; la cuantización de canales sobresale en ambos frentes.

Valores: donde la intuición se rompe

La caché de valores no muestra un patrón de canal-atípico. Parece ser bastante plano. Por sí solo, por el argumento de rango, podríamos esperar que cualquiera de estos ejes produzca una calidad similar de compresión.

No lo hacen. Independientemente de cómo se implemente la gestión de claves (los resultados de 2,80 y 2,88), la compresión por canal de valores colapsa el modelo.

Y aquí está la cuestión: si se mide esta pérdida utilizando el error de reconstrucción bruto en el tensor original para el que cada valor se comprimió, la cuantización por canal de valores en realidad parece ligeramente mejor, con 3,73 frente a 4,57. Si se validó la compresión de la manera obvia, se elegiría la configuración que destruye el modelo.

Error de cuantización de caché de valores en Llama-2-13B, medido de dos maneras. La métrica del tensor almacenado y la métrica de la salida consumida no coinciden en más de una orden de magnitud.

La resolución es que la caché de valores nunca se lee directamente. Se consume mediante un producto matricial: la salida de atención es una suma ponderada de vectores de valores a lo largo de los tokens, con puntuaciones de atención softmax como pesos. Debido a esto, el error relevante es el que se introduce durante este proceso y no dentro de los propios tensores. Medido en términos de la salida de atención, el orden se invirtió por completo. El error relativo informado por KIVI para la salida de atención debido a la cuantización de vectores de valores por token fue de 3,55 en comparación con 49,89 para la cuantización por canal, más de catorce veces mayor para lo que parecía la mejor opción basada en cómo se comprimió.

La explicación es la escasez de atención, que midieron en un 84,3 por ciento. La mayoría de la información contenida en la salida se puede atribuir a unos pocos tokens muy importantes. La cuantización por token confina el error de cada token a ese token, por lo que los errores en tokens no importantes se multiplican por pesos de atención casi cero y efectivamente desaparecen. La cuantización por canal esparce el error de cada token a lo largo de una escala de canal compartida, por lo que los tokens mal representados contaminan la representación de los que importan. La escasez que hace que la atención sea eficiente es la misma propiedad que hace que la cuantización por token sea segura.

La lección transferible es más amplia que la caché KV: medir el error de compresión donde se consume el tensor, no donde se almacena. Un supuesto implícito hecho por el error de reconstrucción es que cada componente de un tensor tiene un peso igual cuando contribuye a la salida final. La atención explícitamente no. Cualquier operación descendente que pondera, bloquea o espacia su entrada rompe ese supuesto. Los lectores familiarizados con mi artículo anterior sobre ceguera a las contradicciones en los sistemas de recuperación reconocerán que estos resultados son similares a las fallas descritas anteriormente: métricas fáciles de calcular que informan sobre algo diferente a lo que se pretendía.

Las incrustaciones rotativas complican las claves

Hay algunos problemas con el uso de incrustaciones de posición rotativas (RoPE). RoPE gira pares de canales en función de la posición relativa de cada token. Ese mezclado disuelve parcialmente la estructura de canal fijo que hizo que la cuantización de claves por canal funcionara en primer lugar, un canal atípico se gira hacia sus vecinos y los vecinos heredan el rango. La respuesta de KVQuant es el orden: cuantizar claves antes de aplicar la rotación y aplicar RoPE después de la descuantización. Junto con la cuantización de claves por canal, tipos de datos no uniformes y aislamiento de una pequeña fracción de valores atípicos, esto obtiene una degradación de perplexidad inferior a 0,1 en 3 bits y permite servir LLaMA-7B hasta 1 millón de tokens de contexto en una sola A100-80GB.

También es importante entender el nivel de impacto de RoPE. Los autores del artículo “RotateKV” informaron un aumento del 145 por ciento en los errores de cuantización una vez que se agregó RoPE y notaron que los canales atípicos difieren en las cabezas de atención, lo que explica por qué aplicar una matriz de rotación compartida en todas partes es insuficiente y las rotaciones adaptadas a la cabeza funcionan mejor.

El impuesto de sistemas y por qué no es un detalle

La cuantización por token se adapta bien a la decodificación. Cada token llega; se cuantiza, se agrega a la secuencia (a lo largo de la dimensión del token), nada más se mueve.

Sin embargo, la cuantización por canal no encaja. Como las estadísticas de un canal abarcan tokens que aún no se han generado, no se puede calcular un factor de escala cuando llega un token. El trabajo de KIVI es mantener los tokens más recientes, hasta 128, en precisión completa en un búfer residual y cuantizar en grupos una vez que se hayan acumulado suficientes.

Resulta que el búfer residual se convierte en una carga, en lugar de ser solo algo incidental. En GSM8K con Llama-2-7B, las puntuaciones de precisión completa son de 13,50. Cuantizado completamente a 2 bits con los ejes correctos, la puntuación es de 5,76. Los mismos ejes y los mismos bits, más el búfer residual de tokens recientemente producidos a precisión completa, obtienen una puntuación de 12,74. Una ventana deslizante de tokens recientemente producidos a precisión completa recuperará gran parte de lo que se perdió debido a la cuantización agresiva en problemas de varios pasos difíciles, lo que tendría sentido si consideramos qué tokens estaban siendo atendidos por una cadena de operaciones aritméticas.

Hay un beneficio significativo al hacer todas estas cosas correctamente, como informa KIVI, 2,6 veces menos uso de memoria pico para Llama-2-7B, lo que permite tamaños de lote hasta 4 veces más grandes, así como un rendimiento 2,35 a 3,47 veces mejor en una tarea de servicio real.

Qué hacer con esto

1. Nunca use un cuantizador para ambos. Use cuantizadores diferentes para claves (por canal) y para valores (por token). Un pipeline que aplica un solo cuantizador a “la caché KV” probablemente ya ha sacrificado la mayoría de la calidad posible al usar un pequeño número de bits para representar cada valor.
2. Cuantice las claves antes de RoPE. Esto es una cuestión de corrección en lugar de una cuestión de preferencia.
3. Almacene una ventana de precisión completa de tokens recientemente generados. Aunque almacenar dicha ventana toma muy poca memoria en comparación con lo grande que puede ser la caché, es precisamente esta área la que genera mucha precisión para tareas difíciles.
4. No valide en el error de reconstrucción. Siempre valide en función de la salida de atención o del rendimiento de la tarea final. La métrica de almacenamiento no es solo ruidosa, para los valores apunta en la dirección equivocada.
5. No valide en pruebas de opción múltiple de contexto corto. Los autores de KIVI evitan deliberadamente tareas cerradas como MMLU para esta evaluación, porque un solo paso de decodificación que lee los logits de salida apenas ejercita la caché en absoluto. Cualquier evaluación que no construya una caché con el tiempo y luego realice generación a partir de ella nunca podrá observar los fallos inherentes al diseño del sistema.

Hacia dónde se dirige el trabajo

Aunque todavía hay algo que hacer con respecto a la naturaleza geométrica del problema, muchos investigadores continúan estudiando la forma en que los canales atípicos se distribuyen entre las diversas cabezas del transformador y cómo las limitaciones del hardware afectan qué agrupaciones son más baratas: InnerQ pliega la normalización de claves por canal en los pesos de claves y consultas durante el prefijo. Por lo tanto, no se incurre en sobrecarga adicional en tiempo de ejecución. Además, InnerQ almacena ventanas de alta precisión para tokens recientemente generados y tokens de sumidero de atención. Al hacerlo, InnerQ elimina la oportunidad de que los valores atípicos en el canal de sumidero contaminen los canales vecinos.

Otros proponen que, en lugar de almacenar toda la caché, deberíamos almacenar solo suficiente información para poder rematerializar la clave y/o el valor (s) a demanda a partir de una representación de caché más pequeña.

Finalmente, es importante recordar que la precisión no es el único parámetro que la cuantización afecta. Una investigación publicada recientemente demostró la degradación de la alineación resultante de la cuantización de las cachés KV. Además, esta investigación documentó la degradación de la alineación incluso en entornos de servicio de vLLM de producción que utilizan cachés FP8 junto con un protocolo de recuperación sin entrenamiento que restauró hasta el 97 por ciento de lo que se perdió en términos de alineación. Así, aunque una configuración puede mantener sus resultados de referencia, no necesariamente significa que retiene todos los demás parámetros relevantes que le importan.

El principio general

La idea de cuantización se ha enmarcado como un “presupuesto de precisión”: ¿cuántos bits puedo permitirme sacrificar? La caché KV muestra que la pregunta más útil es estructural. La precisión se asigna en grupos; el grupo es la unidad de daño, y el eje que se elige para agrupar determina qué elementos comparten su destino. El eje correcto es aquel en el que se consume su tensor, es decir, la forma en que se utiliza su tensor y NO cómo aparece su tensor cuando se almacena en memoria. Las claves se utilizan a través de un cálculo de producto escalar contra la consulta. Un canal corrupto puede envenenar todas las puntuaciones. Los valores se consumen a través de un cálculo de promedio ponderado espaciado a lo largo de los tokens. Por lo tanto, un token corrupto simplemente se pesa. Dos tensores de dimensiones idénticas y generados por dos capas consecutivas se tratan de manera diferente. Vale la pena preguntar sobre cualquier activación que se planea comprimir: ¿qué operación contrata esto, y ¿mi agrupación respeta eso?

Dos tensores de dimensiones idénticas y generados por dos capas consecutivas se tratan de manera diferente. Es importante preguntar sobre cualquier activación que se planea comprimir: ¿qué operación contrata esto, y ¿mi agrupación respeta eso? La ira de cómputo a través de los tokens. Por lo tanto, un token corrupto simplemente se pesa. Dos tensores de dimensiones idénticas y generados por dos capas consecutivas se tratan de manera diferente. Es importante preguntar sobre cualquier activación que se planea comprimir: ¿qué operación contrata esto, y ¿mi agrupación respeta eso?