Comment la Quantization Aware Training permet de récupérer la précision en inferance à faible précision

TL;DR

• L’entraînement quantifié (QAT) et la distillation quantifiée (QAD) étendent la quantification post-entraînement (PTQ) en entraînant le modèle pour s’adapter à la basse précision, entraînant souvent une récupération de précision lors de l’inférence.
• QAT utilise une quantification fictive dans le parcours avant et ne nécessite pas de support matériel natif ; il peut s’intégrer dans des pipelines d’entraînement existants et utilise des techniques comme l’estimation par passage direct (STE).
• QAD ajoute une étape de distillation entre un étudiant quantifié et un professeur en haute précision pour mieux compenser les erreurs de quantification ; PTQ et QAT/QAD fonctionnent avec les formats du TensorRT Model Optimizer tels que NVFP4 et MXFP4.
• Des résultats réels montrent des gains significatifs sur certains modèles et benchmarks (par ex., Llama Nemotron Super) avec QAD ; de nombreux modèles conservent une précision élevée avec PTQ seul ; la qualité des données et les hyperparamètres restent cruciaux.
• NVIDIA fournit des API expérimentales dans le Model Optimizer pour QAT et QAD, avec des guides et tutoriels pour préparer et déployer des modèles en faible précision sur les GPUs NVIDIA.

Contexte et arrière-plan

La quantification post-entraînement (PTQ) est une approche courante pour compresser et déployer des modèles en appliquant un échelonnement numérique vers des types de données de faible précision. Cependant, PTQ peut ne pas suffire pour certaines architectures ou charges de travail, ce qui conduit à l’exploration de Quantization Aware Training (QAT) et Quantization Aware Distillation (QAD). QAT vise à préserver la précision en entraînant le modèle pour fonctionner sous arithmétique de faible précision pendant une phase d’entraînement supplémentaire après le pré-entraînement. Contrairement à PTQ, qui quantifie après un entraînement complet en utilisant un jeu de calibration, QAT entraîne le modèle avec des valeurs quantisées dans le chemin avant. En pratique, le flux QAT est quasi identique à PTQ, la différence clé résidant dans la boucle d’entraînement qui suit la recette de quantisation. QAT utilise typiquement des poids et activations « faux quantisés » dans le passage avant. La faible précision est représentée dans un type de données de plus haute précision par un opérateur de quantisation/desquantisation, si bien qu’un support matériel natif n’est pas nécessaire. Par exemple, QAT pour NVFP4 peut être réalisé sur des GPU Hopper via quantification simulée, et les gradients de rétropropagation sont calculés en haute précision via l’estimation par passage direct (STE). QAD, quant à lui, introduit une composante distillation. L’objectif est de récupérer la précision après la quantification post-entraînement en alignant les sorties d’un étudiant quantifié avec celles d’un professeur en pleine précision entraîné sur les mêmes données. Le professeur demeure en précision complète, tandis que les calculs de l’étudiant sont quantifiés durant la distillation. La perte de distillation mesure l’écart entre les prédictions quantifiées et le comportement du professeur, guidant les poids et activations de faible précision vers le comportement du professeur. Le TensorRT Model Optimizer prend en charge les mêmes formats de quantisation que le flux PTQ, y compris FP8, NVFP4, MXFP4, INT8 et INT4. La sortie du QAT dans l’Optimiseur est un nouveau modèle en précision d’origine avec des poids mis à jour et des métadonnées critiques pour la conversion vers le format cible. Le flux de travail QAT est pratiquement identique à PTQ jusqu’à la boucle d’entraînement, et les développeurs peuvent exploiter des composants d’entraînement familiers comme les optimiseurs, les planificateurs et les stratégies de taux d’apprentissage. Pour les développeurs, cette approche s’intègre bien dans les flux Hugging Face et PyTorch, permettant d’expérimenter les formats de faible précision dans des pipelines établis. L’implémentation actuelle du Model Optimizer considère QAD comme expérimental. Le processus commence par une étape similaire à QAT et PTQ pour le modèle étudiant, puis une configuration de distillation définit des éléments tels que le modèle professeur, les arguments d’entraînement et la fonction de perte de distillation. NVIDIA précise que les API QAD sont expérimentales et feront l’objet d’améliorations pour simplifier l’application. Pour ceux qui souhaitent une pratique guidée, le billet de blog pointe vers un notebook Jupyter et la section QAD du dépôt du Model Optimizer pour les exemples et la documentation les plus récents. Tous les modèles n’exigent pas QAT ou QAD — beaucoup conservent plus de 99,5% de leur précision initiale sur les benchmarks clés avec PTQ seul. Cependant, dans des cas comme le Llama Nemotron Super, QAD peut apporter des gains notables. Sur des benchmarks tels que GPQA Diamond, LIVECODEBENCH, Math-500 et AIME 2024, on observe une récupération de 4 à 22% de précision via QAD, à l’exception du LIVECODEBENCH. L’efficacité de QAT et QAD dépend fortement de la qualité des données d’entraînement, des hyperparamètres et de l’architecture du modèle. Lorsque l’on quantifie en 4 bits, les formats NVFP4 bénéficient de facteurs d’échelle plus granulaires et d’une plus grande précision, ce qui aide à préserver les signaux faibles et les valeurs aberrantes. Dans le cadre du OpenVLM sur le Leaderboard Hugging Face, l’écart de performance entre les meilleurs modèles et les modèles moins bien classés est souvent minime, soulignant l’importance du choix de la quantification selon la tâche et le modèle. Les techniques QAT et QAD visent à équilibrer les avantages de l’exécution FP4 et la robustesse de l’entraînement à haute précision.

Pourquoi c’est important (impact pour les développeurs/entreprises)

Pour les développeurs et les entreprises qui déploient l’IA à grande échelle, pouvoir exécuter l’inférence de manière efficace en faible précision se traduit par une vitesse plus élevée, une utilisation mémoire réduite et des coûts allégés. PTQ offre une réduction et une vitesse de base, mais QAT et QAD fournissent des voies pour récupérer ou améliorer la précision lorsque la calibration simple échoue. L’analyse NVIDIA montre que, selon le modèle et la tâche, QAD peut combler l’écart entre l’inférence en faible précision et les baselines haute précision de plusieurs points de pourcentage, ce qui peut être significatif en production. L’idée pratique est d’avoir une stratégie de quantisation complète: PTQ pour une compression large; QAT lorsque l’on a besoin d’une plus grande précision d’inférence en FP4 et que le temps d’entraînement est acceptable; et QAD lorsque l’étape supplémentaire de distillation peut apporter des gains sans changer l’architecture sous-jacente. La compatibilité PyTorch et Hugging Face via le TensorRT Model Optimizer réduit l’obstacle à l’expérimentation, permettant aux équipes de tester des formats comme NVFP4 et MXFP4 dans des flux de travail familiers. Pour les déploiements réels, cela signifie choisir le bon format de quantification pour chaque charge de travail, appliquer la bonne durée de QAT ou QAD et valider les performances avec des données représentatives.

Détails techniques ou Mise en œuvre

• Aperçu de QAT: entraîne le modèle avec des valeurs quantisées dans le chemin avant pour simuler le comportement de déploiement et exposer les effets d’arrondi et de clipping pendant l’entraînement. L’approche utilise typiquement un opérateur de quantification/desquantification avec des gradients calculés en haute précision via STE. Cela rend QAT faisable sans support matériel natif et s’intègre aux pipelines d’entraînement existants.
• Boucle d’entraînement et durée: QAT nécessite une boucle d’entraînement avec des paramètres tels que optimiseur et schéma de taux d’apprentissage. NVIDIA recommande environ 10% des epochs initiaux pour QAT, bien que dans certains domaines de LLM, des améliorations soient observées avec moins de 1% du pré-entraînement.
• Formats et métadonnées: QAT dans Model Optimizer prend en charge les formats FP8, NVFP4, MXFP4, INT8 et INT4. La sortie est un modèle en précision d’origine avec des poids mis à jour et des métadonnées critiques pour la conversion vers le format cible.
• Détails de QAD: QAD suit les étapes initiales similaires à QAT et PTQ pour le modèle étudiant, puis introduit une configuration de distillation. L’étudiant quantifié effectue les calculs pendant la distillation, pendant que le professeur reste en précision complète. La perte de distillation aligne les sorties de l’étudiant sur le comportement du professeur et guide l’ajustement des poids et activations sous quantisation. NVIDIA précise que les API QAD sont expérimentales et seront améliorées pour simplifier l’utilisation.
• L’intégration et le flux de travail: le TensorRT Model Optimizer propose des API compatibles avec Hugging Face et PyTorch pour faciliter la préparation de QAT et QAD, tout en tirant parti des flux de travail familiers. L’objectif est de permettre aux équipes de compresser, affiner et déployer des modèles sur les GPUs NVIDIA avec un équilibre entre vitesse, taille et précision.

Points clés

• QAT et QAD étendent PTQ en introduisant un apprentissage spécifique pour gérer l’arithmétique de faible précision, aboutissant souvent à une précision d’inférence plus élevée.
• QAT utilise une quantification fictive dans le passage avant et ne nécessite pas de support matériel natif; il est compatible avec FP8, NVFP4, MXFP4, INT8 et INT4.
• QAD ajoute une étape de distillation pour rapprocher les sorties d’un étudiant quantifié de celles d’un professeur en haute précision, compensant les erreurs de quantification pendant l’entraînement.
• PTQ demeure efficace pour de nombreux modèles, mais QAT/QAD peuvent offrir des gains notables selon le modèle et la tâche, notamment à 4 bits.
• La qualité des données, les hyperparamètres et l’architecture du modèle influencent fortement l’efficacité des méthodes de quantisation; les expériences sont soutenues par les API du TensorRT Model Optimizer et l’écosystème PyTorch/Hugging Face.

Références

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