Cómo la quantización consciente (QAT) recupera precisión en inferencia de baja precisión

TL;DR

• El entrenamiento con cuantización consciente (QAT) y la distilación cuantizada (QAD) extienden la cuantización posterior al entrenamiento (PTQ) al entrenar el modelo para adaptarse a la baja precisión, a menudo recuperando mayor precisión durante la inferencia.
• QAT utiliza cuantización falsa en la pasada hacia adelante y no requiere soporte de hardware nativo; puede integrarse en flujos de entrenamiento existentes y usa técnicas como estimación por paso directo (STE).
• QAD añade una componente de distilación entre un estudiante cuantizado y un profesor de alta precisión para compensar errores de cuantificación; PTQ y QAT/QAD funcionan con formatos del TensorRT Model Optimizer como NVFP4 y MXFP4.
• Resultados del mundo real muestran mejoras en ciertos modelos y benchmarks (p. ej., Llama Nemotron Super) al aplicar QAD; muchos modelos mantienen alta precisión con PTQ; la calidad de los datos y los hiperparámetros siguen siendo clave.
• NVIDIA ofrece API experimentales en el Model Optimizer para QAT y QAD, con guías y tutoriales para ayudar a preparar y desplegar modelos en baja precisión en GPUs de NVIDIA.

Contexto y antecedentes

La cuantización post-entrenamiento (PTQ) es un enfoque común para comprimir y desplegar modelos aplicando escalamiento numérico a pesos en tipos de datos de menor precisión. Sin embargo, PTQ puede no ser suficiente para algunas arquitecturas o cargas de trabajo, lo que lleva a la exploración de Quantization Aware Training (QAT) y Quantization Aware Distillation (QAD). QAT busca preservar la precisión entrenando el modelo para operar con aritmética de baja precisión durante una fase adicional de entrenamiento tras el preentrenamiento. A diferencia de PTQ, que cuantiza después del entrenamiento completo usando un conjunto de calibración, QAT entrena el modelo con valores cuantizados en la pasada hacia adelante. En la práctica, el flujo de QAT es casi idéntico a PTQ, la diferencia clave es la inclusión de la fase de entrenamiento tras aplicar la receta de cuantización. QAT suele funcionar con pesos y activaciones “falsamente cuantizados” en la pasada hacia adelante. La menor precisión se representa dentro de un tipo de datos de mayor precisión mediante un operador de cuantización/descuantización, por lo que no se requiere soporte de hardware nativo. Por ejemplo, QAT para NVFP4 puede ejecutarse en GPUs Hopper usando cuantización simulada, y los gradientes de retropropagación se calculan en precisión más alta mediante estimación por paso directo (STE). QAD, por otro lado, introduce una componente de distilación. El objetivo es recuperar precisión tras la cuantización post-entrenamiento alineando las salidas del estudiante cuantizado con las del modelo maestro en alta precisión, entrenado con los mismos datos. El maestro permanece en precisión total, mientras que los cálculos del estudiante se realizan con cuantización durante la distilación. La pérdida de distilación mide cuánto se desvía la predicción cuantizada y guía los pesos y activaciones de baja precisión para ajustarse al comportamiento del maestro. El TensorRT Model Optimizer admite los mismos formatos de cuantización que el flujo PTQ, incluyendo FP8, NVFP4, MXFP4, INT8 e INT4. La salida del QAT en Model Optimizer es un modelo en la precisión original con pesos actualizados y metadatos críticos para convertirlo al formato objetivo. El flujo de trabajo de QAT es prácticamente idéntico a PTQ hasta la etapa de entrenamiento, y los desarrolladores pueden aprovechar componentes familiares de entrenamiento como optimizadores, programadores y estrategias de tasa de aprendizaje. Para los desarrolladores, este enfoque se integra bien en flujos de trabajo de PyTorch y Hugging Face, permitiendo experimentar con formatos de baja precisión dentro de pipelines familiares. La implementación actual de QAD en Model Optimizer es experimental. Comienza con un proceso similar a QAT y PTQ para el modelo estudiante, seguido por una configuración de distilación que define elementos como el modelo maestro, argumentos de entrenamiento y la función de pérdida de distilación. NVIDIA señala que las API de QAD continuarán mejorándose para simplificar su uso. Para quienes buscan guía práctica, el blog señala un cuaderno Jupyter y la sección QAD del repositorio de Model Optimizer para ejemplos y documentación actualizados. No todos los modelos requieren QAT o QAD: muchos conservan más del 99.5% de su precisión original en benchmarks clave solo con PTQ. Aun así, en casos como Llama Nemotron Super, QAD puede generar mejoras notables. En benchmarks como GPQA Diamond, LIVECODEBENCH, Math-500 y AIME 2024, se observa recuperación de 4–22% de precisión con QAD, con excepción de LIVECODEBENCH. La efectividad de QAT y QAD depende fuertemente de la calidad de los datos de entrenamiento, de los hiperparámetros elegidos y de la arquitectura del modelo. Al cuantificar a 4 bits, formatos como NVFP4 se benefician de factores de escala más granulares. Desde una perspectiva práctica, NVFP4 tiende a ofrecer pequeñas pero significativas mejoras frente a MXFP4 en ciertos escenarios, destacando la importancia de la selección de formato para una carga de trabajo dada. En el contexto de OpenVLM y el Leaderboard de Hugging Face, la brecha de rendimiento entre el mejor modelo y los que están menos bien posicionados suele ser pequeña, lo que subraya que las elecciones de cuantización afectan de manera sensible los resultados. En conjunto, QAT y QAD amplían los beneficios de PTQ al enseñar a los modelos a adaptarse directamente a entornos de baja precisión, logrando recuperaciones de precisión donde la calibración simple no alcanza. Como muestran benchmarks como Math-500 y AIME 2024, estas técnicas pueden cerrar la brecha entre la inferencia de baja precisión y las referencias de alta precisión, combinando la eficiencia de la ejecución FP4 con la robustez del entrenamiento de alta precisión. La calidad de los datos de entrenamiento, los hiperparametros y la arquitectura del modelo siguen siendo determinantes para el éxito de estas técnicas. Con TensorRT Model Optimizer, estos flujos de trabajo avanzados están disponibles a través de APIs familiares en PyTorch y Hugging Face, facilitando la experimentación con formatos como NVFP4 y MXFP4.

¿Por qué es importante (impacto para desarrolladores/empresas)?

Para equipos que implementan IA a gran escala, la capacidad de ejecutar inferencia de forma eficiente en baja precisión se traduce en mayor rendimiento, menor uso de memoria y costos reducidos. PTQ ofrece una base de compresión y velocidad, pero QAT y QAD proporcionan caminos para recuperar o mejorar la precisión cuando la calibración simple no es suficiente. Los análisis de NVIDIA muestran que, dependiendo del modelo y la tarea, QAD puede cerrar la brecha entre la inferencia de baja precisión y las baselines de alta precisión en varios puntos porcentuales, lo que puede ser significativo en producción. La idea práctica es disponer de una estrategia de cuantización completa: PTQ para una compresión amplia; QAT cuando se necesita mayor precisión de inferencia en FP4 y el tiempo de entrenamiento es aceptable; y QAD cuando la distilación adicional aporta mejoras sin cambiar la arquitectura subyacente. La interoperabilidad con PyTorch y Hugging Face a través del TensorRT Model Optimizer facilita la experimentación, permitiendo a los equipos probar formatos como NVFP4 y MXFP4 dentro de flujos de trabajo familiares. Para despliegues reales, esto significa seleccionar el formato de cuantización adecuado para cada carga de trabajo, aplicar la cantidad adecuada de QAT o QAD y validar el rendimiento con datos representativos.

Detalles técnicos o Implementación

• Panorama general de QAT: entrena el modelo con valores cuantizados en el camino hacia adelante para simular el comportamiento de despliegue y exponer efectos de redondeo y recorte durante el entrenamiento. Se utiliza típicamente un operador de cuantización/descuantización con gradientes calculados en alta precisión mediante STE. Esto hace que QAT sea factible sin soporte de hardware nativo e se integra a pipelines de entrenamiento ya establecidos.
• Bucle de entrenamiento y duración: se recomienda dedicar aproximadamente el 10% de los epochs iniciales al QAT, con beneficios en LLM a veces observables incluso con menos del 1% del preentrenamiento.
• Formatos y metadatos: QAT en Model Optimizer admite FP8, NVFP4, MXFP4, INT8 e INT4. La salida es un modelo en la precisión original con pesos actualizados y metadatos para convertirlo al formato objetivo.
• Detalles de QAD: se inicia con un proceso similar a QAT y PTQ para el modelo estudiante, seguido de una configuración de distilación que especifica el modelo maestro, argumentos de entrenamiento y la función de pérdida de distilación. El estudiante cuantizado realiza las operaciones durante la distilación y el maestro permanece en plena precisión. NVIDIA señala que las API de QAD son experimentales y se mejorarán para simplificar su uso.
• Integración y flujo de trabajo: el TensorRT Model Optimizer ofrece API compatibles con Hugging Face y PyTorch para facilitar la preparación de QAT y QAD y aprovechar flujos de trabajo familiares. El objetivo es permitir comprimir, ajustar y desplegar modelos en GPUs NVIDIA con un equilibrio entre velocidad, tamaño y precisión.

Conclusiones clave

• QAT y QAD amplían PTQ al introducir aprendizaje específico para manejar la aritmética de baja precisión, con frecuencia resultando en mayor precisión de inferencia.
• QAT utiliza cuantización falsa en la pasada hacia adelante y no requiere soporte de hardware nativo; es compatible con FP8, NVFP4, MXFP4, INT8 e INT4.
• QAD añade una etapa de distilación para alinear las salidas de un estudiante cuantizado con las del profesor de alta precisión, compensando errores de cuantificación durante el entrenamiento.
• PTQ sigue siendo eficaz para muchos modelos, pero QAT/QAD pueden ofrecer mejoras significativas según el modelo y la tarea, especialmente al cuantizar a 4 bits.
• La calidad de los datos de entrenamiento, los hiperparámetros y la arquitectura de modelado influyen fuertemente en la efectividad de estas técnicas; las pruebas se apoyan en las APIs experimentales del TensorRT Model Optimizer con integración a PyTorch/Hugging Face.

Referencias

• https://developer.nvidia.com/blog/how-quantization-aware-training-enables-low-precision-accuracy-recovery/