OneComp: One-Line Revolution for Generative AI Model Compression

Resumen Técnico: OneComp

1. El Problema

La implementación de modelos fundacionales (Foundation Models) de IA generativa se enfrenta a cuellos de botella críticos relacionados con el huella de memoria, la latencia y los costos de hardware. Los modelos con decenas o cientos de miles de millones de parámetros exceden la capacidad de memoria de las GPUs comerciales, haciendo prohibitiva su ejecución en precisión completa.

Aunque la cuantización post-entrenamiento (PTQ) es una estrategia clave para mitigar estos problemas reduciendo la precisión de los pesos sin reentrenar, su adopción práctica es difícil debido a:

• Un ecosistema fragmentado de algoritmos de cuantización.
• La necesidad de configuraciones manuales complejas (presupuestos de precisión, estrategias de calibración, asignación de bits).
• La falta de un flujo de trabajo unificado que se adapte automáticamente a las restricciones de hardware disponibles.
• La brecha entre la investigación teórica (que a menudo requiere recursos masivos) y la implementación en producción.

2. Metodología

OneComp es un framework de código abierto diseñado para transformar el flujo de trabajo de compresión en una tubería (pipeline) reproducible y adaptativa a los recursos. Su diseño se basa en tres principios fundamentales:

• Flujo de Trabajo End-to-End Automatizado: Dado un identificador de modelo y el hardware disponible, OneComp inspecciona automáticamente la arquitectura, planifica la asignación de bits y ejecuta las etapas de cuantización sin intervención manual.
• Arquitectura de Refinamiento Basada en "Pivote": Trata el primer checkpoint cuantizado (obtenido mediante PTQ por capas) como un punto de partida (pivote) desplegable. Las etapas posteriores (refinamiento por bloques, global y ajuste fino) mejoran iterativamente este mismo modelo, permitiendo al usuario detenerse cuando se agote el presupuesto de hardware.
• Motor Adaptativo a Recursos: El sistema escala automáticamente desde métodos ligeros (una sola GPU) hasta optimizaciones globales intensivas, según la memoria disponible.

Componentes Técnicos Clave:

1. AutoBit (Planificación de Precisión Mixta):

   • Resuelve un problema de optimización lineal entera para asignar diferentes anchos de bits a diferentes módulos del modelo bajo un presupuesto global de memoria.
   • Utiliza estadísticas de activación (Hessiano diagonal) para identificar capas sensibles y asignarles más bits, evitando la fragilidad de la cuantización uniforme.

2. Preprocesamiento por Transformaciones Equivalentes:

   • Mitiga los "outliers" (valores extremos) en pesos y activaciones mediante transformaciones matemáticas que no alteran la función del modelo.
   • Incluye escalado de canales (SmoothQuant), rotaciones basadas en Hadamard y rotaciones aprendidas (SpinQuant) para redistribuir la energía de los outliers y mejorar la estabilidad de la cuantización.

3. Etapas de Cuantización (Regímenes de PTQ):

   • PTQ por Capas (Layer-wise): Procesa una capa lineal a la vez. OneComp introduce QEP (Corrección de Propagación de Error de Cuantización) para ajustar el objetivo de cuantización de una capa basándose en los errores acumulados de las capas anteriores, y LPCD (Descenso de Coordenadas Consciente de Submódulos) para optimizar conjuntos de capas acopladas (como las proyecciones Q, K, V o las capas MLP) simultáneamente.
   • PTQ por Bloques (Block-wise): Refina un bloque completo de Transformer a la vez. Incluye compensación por bloques (ajustar parámetros FP16 restantes antes de cuantizar) y distilación por bloques (refinar parámetros ya cuantizados minimizando la discrepancia con el modelo maestro).
   • PTQ Global: Optimiza todo el modelo simultáneamente mediante distilación de conocimiento (KL-divergencia) y alineación de estados intermedios, capturando dependencias de largo alcance.

4. Formatos de Representación:

   • JointQ: Optimiza conjuntamente escalas y pesos enteros para rangos de 3-4 bits, superando a GPTQ estándar.
   • DBF y MDBF: Formatos de factores binarios estructurados para rangos extremos (1-2 bits). MDBF mejora a DBF reemplazando el envelope de magnitud de rango único por uno de rango superior ($l$), permitiendo una mayor expresividad en magnitudes.

5. Ajuste Fino (Fine-tuning):

   • Soporte opcional para QLoRA, permitiendo el ajuste de adaptadores de bajo rango sobre el backbone cuantizado para recuperación de errores o adaptación de dominio.

3. Contribuciones Clave

• Unificación de Regímenes: OneComp integra PTQ por capas, por bloques y global en un solo pipeline coherente, eliminando la necesidad de elegir entre herramientas dispares.
• Corrección de Error de Propagación: La introducción de QEP y LPCD permite que la cuantización por capas (la más eficiente en memoria) alcance una calidad cercana a métodos mucho más costosos, corrigiendo sistemáticamente los errores acumulados.
• Planificación Automática de Bits: AutoBit automatiza la asignación heterogénea de bits, demostrando que la cuantización uniforme es subóptima incluso con el mismo costo de memoria total.
• Soporte para Bits Extremos: Implementación de MDBF, que supera las limitaciones de los factores binarios anteriores en el régimen de 1-2 bits, manteniendo precisión utilizable.
• Interfaz de "Una Línea": Permite a usuarios no expertos obtener modelos cuantizados de alta calidad con una sola llamada, mientras ofrece control granular a expertos.

4. Resultados Experimentales

Las evaluaciones se realizaron en modelos LLaMA-3 y Qwen de diversos tamaños:

• Asignación de Bits Mixta (AutoBit): En LLaMA-3-8B, la asignación activa-aware superó significativamente a la uniforme. A 4.16 bits promedio, la versión uniforme tuvo una perplejidad de 9.53, mientras que AutoBit logró 6.85 (cercano al modelo de precisión completa de 6.14). En configuraciones de 4 bits por canal, la cuantización uniforme colapsó (PPL > 300), mientras que AutoBit mantuvo un PPL de 7.04.
• Corrección de Error (QEP/LPCD): En cuantización de 4 bits por capa, QEP redujo la perplejidad de 12.66 (GPTQ base) a 6.66, recuperando casi toda la precisión del modelo original. LPCD ofreció mejoras adicionales en configuraciones agresivas (3 bits).
• JointQ vs. GPTQ: JointQ superó consistentemente a GPTQ estándar en precisión y perplejidad, especialmente en configuraciones de 3 bits y por canal, demostrando que la optimización conjunta de escalas e enteros es superior a la optimización alternada.
• Cuantización Extrema (1-2 bits): MDBF superó a DBF en todos los modelos probados. En LLaMA-3-8B a 1.00 bits por peso, MDBF logró una precisión (ACC) de 0.518 frente a 0.505 de DBF, y una perplejidad significativamente menor.
• Mejora Monótona: Los experimentos mostraron que la calidad del modelo mejora monótonamente a medida que se invierten más recursos computacionales (pasando de PTQ por capas -> por bloques -> global), validando la arquitectura de refinamiento progresivo.

5. Significado e Impacto

OneComp representa un cambio de paradigma en la implementación de modelos de IA generativa:

• Democratización de la Compresión: Reduce la barrera de entrada para la compresión de modelos, permitiendo que organizaciones con recursos limitados desplieguen modelos grandes en hardware estándar.
• Puente entre Investigación y Producción: Convierte algoritmos de investigación complejos en un flujo de trabajo automatizado y robusto, acelerando la adopción de técnicas de vanguardia.
• Eficiencia Económica: Al permitir el despliegue de modelos de 70B+ parámetros en una sola GPU o reducir drásticamente los costos de inferencia, OneComp aborda directamente la sostenibilidad económica de la IA generativa.
• Visión de Código Abierto: Al liberar el framework como open-source, busca establecer un estándar comunitario para la compresión de modelos, fomentando la innovación colaborativa en algoritmos y hardware.

En resumen, OneComp no es solo una herramienta de cuantización, sino un ecosistema completo que gestiona la complejidad de la compresión de modelos, garantizando que la calidad del modelo sea una función monótona de los recursos invertidos, sin requerir configuración manual experta.