التدريب والتكرير المدركان للكمية: تعزيز دقة الدقة المنخفضة باستخدام مُحسِّن نموذج TensorRT

In the relentless pursuit of efficient AI model deployment, post-training quantization (PTQ) has become a standard technique for compressing models. However, when PTQ falls short in maintaining accuracy at lower precisions, Quantization Aware Training (QAT) and Quantization Aware Distillation (QAD) offer promising alternatives. This blog post dives into QAT and QAD, showcasing how they can be leveraged with the TensorRT Model Optimizer to prepare models for low-precision inference, specifically focusing on workflows compatible with Hugging Face and PyTorch. It also highlights the deployment of these optimized models using TensorRT-LLM.

Understanding Quantization Aware Training (QAT)

QAT is a technique where the model is trained to accommodate low-precision arithmetic after pre-training. Unlike PTQ, which quantizes a fully trained model, QAT simulates quantization during the training process, allowing the model to adapt its weights and activations to the limitations of lower precision data types.

• How it works: QAT uses “fake quantization” during the forward pass, representing lower precision values within a higher data type using quantize/dequantize operators. This allows training on standard hardware without native low-precision support.
• The Goal: The primary goal is high accuracy quantized inference, differentiating it from quantized training focused on training efficiency.
• Implementation: During QAT, rounding and clipping errors are exposed to the loss function, allowing the model to recover and adapt.
• Model Optimizer Support: Model Optimizer supports quantization formats such as FP8, NVFP4, MXFP4, INT8, and INT4. Applying QAT with Model Optimizer requires only a small amount of extra code by adding a training loop to the PTQ workflow. Fine-tuning the model for about 10% of the initial training epochs, or even less for LLMs, is generally sufficient to restore model quality.

Exploring Quantization Aware Distillation (QAD)

QAD aims to recover accuracy post-quantization, while simultaneously performing knowledge distillation. The key difference from traditional knowledge distillation is that the “student” model utilizes fake quantization during its forward pass.

• Teacher-Student Setup: QAD uses a teacher model (the original higher precision model) to guide a student model (the quantized model).
• Distillation Loss: A distillation loss function is used to measure the deviation between the quantized student’s outputs and the full-precision teacher’s outputs.
• The Process: The student’s computations are fake quantized during the distillation, exposing quantization-induced mismatches to the distillation loss. This allows the low-precision weights and activations to adapt to the teacher’s behavior.
• Benefits: This leads to higher accuracy recovery compared to PTQ or QAT alone, without impacting inference performance (as the architecture and precision remain the same).
• Model Optimizer APIs: TensorRT Model Optimizer offers experimental APIs for QAD. The process involves applying a quantization recipe to the student model, defining a distillation configuration (teacher model, training arguments, and distillation loss), and then executing the distillation process.

Real-World Impact and Benchmarking

The effectiveness of QAT and QAD varies across models. Some models retain high accuracy (>99.5%) with PTQ alone. However, models like Llama Nemotron Super have shown significant accuracy gains from QAD (4-22% across benchmarks).

• Data and Hyperparameter Dependency: The success of QAT and QAD depends heavily on training data quality, hyperparameter selection, and model architecture.
• NVFP4 vs MXFP4: For 4-bit quantization, NVFP4 can outperform MXFP4 due to finer-grained scaling factors, particularly in tasks involving fine details, like visual question answering on complex charts and graphics. This translates to a steadier alignment between the visual and text components of the model, reducing rounding/clipping errors.
• Visual Question Answering Examples: Differences between NVFP4 and MXFP4 are more prominent in datasets like InfoVQA, which requires detecting small details, compared to DocVQA, which focuses on structured documents where finding the correct field is paramount.

Conclusion

QAT and QAD enhance PTQ by enabling models to adapt to low-precision environments, recovering accuracy where PTQ falls short. These techniques enable developers to harness the efficiency of FP4 execution with near high-precision robustness. The TensorRT Model Optimizer provides access to these workflows through familiar PyTorch and Hugging Face APIs, simplifying experimentation with formats like NVFP4 and MXFP4. Whether the project requires the speed of PTQ, the resilience of QAT, or the accuracy gains of QAD, there is a comprehensive toolkit available to compress, fine-tune, and deploy AI models on NVIDIA GPUs, leading to faster, smaller, and more accurate AI deployments at scale.

Source: NVIDIA Developer Blog