Resumen Técnico: MosaicQuant

Planteamiento del Problema

Si bien la cuantificación de 4 bits ofrece un potencial significativo para reducir la huella de memoria y acelerar la inferencia de los Modelos de Lenguaje de Gran Escala (LLMs), enfrenta una limitación fundamental: el ancho de bits restringido no puede representar fielmente tanto los valores comunes densos (inliers) como los valores de gran magnitud y poco frecuentes (outliers). Esta discrepancia conduce a una degradación sustancial de la precisión.

Las soluciones existentes suelen emplear una cuantificación de precisión mixta, manteniendo los outliers en alta precisión (por ejemplo, FP16 o INT8) mientras se cuantiza el resto a 4 bits. Sin embargo, este enfoque rompe la uniformidad de la ejecución de bajo bit. La necesidad de operaciones de respaldo de alta precisión introduce sobrecarga adicional de movimiento de datos, conversión de formatos y conversión de precisión, lo que en última instancia socava la aceleración de hardware práctica prometida por la inferencia de bajo bit.

Metodología

Los autores proponen MosaicQuant, un paradigma de cuantificación de 4 bits unificado basado en el principio de desagregación de inliers–outliers. En lugar de preservar los outliers en alta precisión, MosaicQuant descompone la matriz de pesos enteramente en componentes de 4 bits: una base densa y un residuo disperso.

1. Desagregación de Inlier–Outlier

El método comienza migrando los outliers de la activación a la matriz de pesos utilizando transformaciones equivalentes (por ejemplo, SmoothQuant o QuaRot). La matriz de pesos migrada $\hat{W}$ se descompone entonces como:
$$\hat{W} \approx Q(\hat{W}) + R$$
Donde:

• Base Densa de 4 bits ($Q(\hat{W})$): Captura fielmente los inliers comunes.
• Residuo Disperso de 4 bits ($R$): Compensa los errores de cuantificación introducidos por los outliers.

Crucialmente, ambos componentes operan en 4 bits, preservando la uniformidad del flujo de ejecución de bajo bit.

2. Asignación de Residuo por Bloques

El artículo identifica que los errores de cuantificación no se distribuyen uniformemente a través de los canales, sino que exhiben una concentración a nivel de bloque de grano fino. Por lo tanto, MosaicQuant asigna el presupuesto del residuo disperso a nivel de bloque en lugar de a nivel de canal.

Para seleccionar los bloques más críticos sin evaluaciones costosas de salida completa, los autores introducen un $\Delta$-score guiado por Hessiano. Esta métrica aproxima la reducción marginal de la distorsión de salida para cada bloque ponderando el error residual con la sensibilidad de la activación (estimación diagonal del Hessiano) de los canales de entrada. El algoritmo selecciona los top-$K$ bloques con los $\Delta$-scores más altos para recibir la compensación del residuo disperso de 4 bits. Los experimentos muestran que un presupuesto disperso de solo el 10% de los bloques es suficiente para recuperar más del 50% de la distorsión de salida.

3. ZipperEngine: Ejecución Unificada

Una representación unificada por sí sola es insuficiente si el residuo disperso se ejecuta mediante un kernel separado, lo que recargaría las activaciones y fusionaría salidas parciales, añadiendo latencia significativa. Para abordar esto, los autores introducen ZipperEngine, un motor de ejecución especializado que:

• Fusiona Kernels: Integra el cómputo del residuo disperso directamente en el kernel GEMM de 4 bits denso.
• Pipeline Solapado: Utiliza un pipeline segmentado y solapado donde la cuantificación, el cómputo denso, la compensación dispersa y las operaciones de reducción-suma se programan concurrentemente a través de Tensor Cores, CUDA Cores y HBM.
• Resultado: Esto elimina la sobrecarga de la programación de kernels separados y el tráfico de memoria redundante, logrando la eficiencia de hardware de un pipeline de inferencia de 4 bits totalmente unificado.

Contribuciones Clave

1. Paradigma MosaicQuant: Un enfoque de cuantificación de 4 bits que realiza una desagregación de inlier–outlier por bloques, mitigando la pérdida de precisión mientras mantiene un patrón de ejecución de bajo bit completamente uniforme.
2. Selección Guiada por Hessiano: Una estrategia de selección de bloques eficiente y consciente de la activación utilizando un $\Delta$-score de Hessiano para identificar los bloques residuales críticos, reemplazando la costosa evaluación exacta de la distorsión de salida.
3. ZipperEngine: Un motor de ejecución de 4 bits denso–disperso que fusiona el cómputo del residuo disperso en el kernel GEMM denso mediante un pipeline solapado, permitiendo aceleraciones prácticas incluso con bloques residuales adicionales.

Resultados Experimentales

Los autores evaluaron MosaicQuant en modelos LLaMA3 (3B, 8B) y Qwen3 (4B, 8B, 14B, 32B).

• Precisión: MosaicQuant preserva una precisión cercana a FP16. Bajo cuantificación W4A4, supera significativamente a otros baselines de bajo bit (RTN, SmoothQuant, QuaRot, SpinQuant, ResQ, Atom). Por ejemplo, en Qwen3-32B, MosaicQuant logra una perplejidad de 8.4, comparado con 9.2 para ResQ y 9.3 para Atom. También mantiene un rendimiento robusto donde otros métodos (como RTN y SmoothQuant) suelen colapsar en modelos Qwen3.
• Throughput: En una NVIDIA RTX 4090, MosaicQuant logra hasta un 1.24× de aceleración sobre el baseline W16A16. Entrega consistentemente el mayor throughput a través de tamaños de batch 4–32 en comparación con TensorRT-LLM y otros baselines de cuantificación.
• Overhead: Las ramas dispersas de 4 bits incurren en un overhead mínimo, con un aumento del uso de memoria de solo 8.5–9.5% y un overhead de latencia que permanece por debajo del 6.2% debido al pipeline fusionado.

Significancia

El artículo afirma que MosaicQuant resuelve el compromiso entre precisión y eficiencia de hardware en la cuantificación de 4 bits de LLMs. Al eliminar la necesidad de operaciones de respaldo de alta precisión y unificar tanto la representación como la ejecución en un único pipeline coherente de bajo bit, MosaicQuant permite el servicio eficiente de LLMs sin sacrificar la regularidad de la ejecución uniforme de 4 bits. Este enfoque demuestra que la compensación dispersa de baja precisión puede recuperar eficazmente los errores inducidos por outliers mientras mantiene las ventajas de throughput de la inferencia densa de bajo bit.