Cuantización Matryoshka

Matryoshka Quantization

February 10, 2025

Autores: Pranav Nair, Puranjay Datta, Jeff Dean, Prateek Jain, Aditya Kusupati

cs.AI

Resumen

La cuantificación de los pesos del modelo es fundamental para reducir los costos de comunicación e inferencia de modelos grandes. Sin embargo, cuantificar modelos, especialmente a bajas precisiones como int4 o int2, implica un compromiso en la calidad del modelo; en particular, se sabe que int2 degrada severamente la calidad del modelo. Por lo tanto, a menudo los profesionales se ven obligados a mantener múltiples modelos con diferentes niveles de cuantificación o a servir un solo modelo que satisfaga mejor el equilibrio entre calidad y latencia. Por otro lado, los tipos de datos enteros, como int8, poseen inherentemente una estructura anidada (Matryoshka) donde enteros de menor tamaño de bits, como int4 o int2, están anidados dentro de los bits más significativos. Este artículo propone la Cuantificación Matryoshka (MatQuant), una novedosa técnica de cuantificación multi-escala que aborda el desafío de necesitar múltiples modelos cuantificados. Permite entrenar y mantener un solo modelo, que luego puede ser servido en diferentes niveles de precisión. Además, debido a la regularización de co-entrenamiento y co-destilación proporcionada por MatQuant, los modelos de precisión int2 extraídos por MatQuant pueden ser hasta un 10% más precisos que la cuantificación estándar int2 (utilizando técnicas como QAT u OmniQuant). Esto representa un progreso significativo en la cuantificación de modelos, demostrado por el hecho de que, con la misma receta, un modelo Gemma-2 9B FFN cuantificado en int2 es más preciso que un modelo Gemma-2 2B FFN cuantificado en int8.

English

Quantizing model weights is critical for reducing the communication and inference costs of large models. However, quantizing models -- especially to low precisions like int4 or int2 -- requires a trade-off in model quality; int2, in particular, is known to severely degrade model quality. Consequently, practitioners are often forced to maintain multiple models with different quantization levels or serve a single model that best satisfies the quality-latency trade-off. On the other hand, integer data types, such as int8, inherently possess a nested (Matryoshka) structure where smaller bit-width integers, like int4 or int2, are nested within the most significant bits. This paper proposes Matryoshka Quantization (MatQuant), a novel multi-scale quantization technique that addresses the challenge of needing multiple quantized models. It allows training and maintaining just one model, which can then be served at different precision levels. Furthermore, due to the co-training and co-distillation regularization provided by MatQuant, the int2 precision models extracted by MatQuant can be up to 10% more accurate than standard int2 quantization (using techniques like QAT or OmniQuant). This represents significant progress in model quantization, demonstrated by the fact that, with the same recipe, an int2 FFN-quantized Gemma-2 9B model is more accurate than an int8 FFN-quantized Gemma-2 2B model.