Quantization Aware Training (QAT) vs. Post-Training Quantization (PTQ)

Jaideep Ray

Follow

8 min read

·

Dec 26, 2024

50

2

Listen

Share

Smaller models => Faster inference => Better outcomes

Introduction

Quantization is a technique for deploying large neural networks efficiently, enabling faster inference and reduced memory footprint.

It involves converting the model’s weights and activations from higher precision (e.g., fp32, fp16, bf16) to lower precision (e.g., int8 or lower).

Two primary quantization approaches exist: Quantization Aware Training (QAT) and Post-Training Quantization (PTQ). Here, we dive into the tradeoffs of using them for LLM and Rec-sys models.

Press enter or click to view image in full size

tl;dr:

Press enter or click to view image in full size

Understanding Quantization Aware Training (QAT)

QAT integrates the quantization process into the model training phase. During training, the model simulates quantization on each forward and backward pass, helping the model learn how to deal with quantization errors in real-time.

Both weights and activations are fake quantized using a specific scheme (int8), and a dequantization step is performed to recover the full-precision values for gradient computation. “fake quantized” means they are transformed as if they were being quantized, but kept in the original data type (e.g. bfloat16) without being actually cast to lower bit-widths. Thus, fake quantization allows the model to adjust for quantization noise when updating the weights, hence the training process is “aware” that the model will ultimately be quantized after training.

The loss function and optimization process are adjusted to account for the quantization effects, ensuring the model learns to minimize the impact of quantization on its performance 5.

Here’s a breakdown of how it works:

1. Forward Pass Simulation: In every forward pass during training, the model simulates the quantization of both weights and activations. This means that the model experiences the effects of reduced precision as if it were already quantized . This simulation introduces quantization errors, which are incorporated into the model’s loss function.
2. Backward Pass with Full Precision: While the forward pass simulates quantization, the backward pass, where gradients are calculated and weights are updated, remains in full precision . This ensures that the model learns to minimize the impact of quantization errors without compromising the accuracy of gradient calculations.
3. Learning and Adaptation: By experiencing quantization errors during training, the model learns to adjust its parameters to mitigate the negative effects of reduced precision . This adaptation makes the model more robust to quantization, leading to better accuracy when the model is actually quantized for deployment.
4. Quantization-Aware Optimization: The optimizer used during training also plays a role in QAT. It considers the quantization errors (due to fake quantize operation) apart from the prediction loss and adjusts the model’s parameters accordingly to minimize the overall loss.

Torchtune (post training pytorch ecosystem library ) includes essential QAT components:

• FakeQuantize: Modules to simulate quantization during training.
• Observer: Modules to collect statistics for quantization parameter calculation (scale and zero-point).
• prepare_qat: Function to prepare a model for QAT by inserting FakeQuantize and Observer modules.
• convert: Function to convert a QAT-trained model to a quantized inference model.

Targeting Specific Layers: QAT is typically most beneficial when applied to nn.Linear layers, which are prevalent in LLMs and sensitive to other post training quantization techniques. Torchtune offers configuration options to specify which parts of the model you want to apply QAT to, often focusing on these linear layers.

It supports token-wise, tensor-wise and channel-wise quantization to ensure flexibility.

Press enter or click to view image in full size

Representative datasets:

• If you have a very large training dataset, performing QAT on the entire dataset for every epoch can significantly increase training time.
• Using a smaller, representative calibration dataset during the initial QAT phase can be more efficient. You are essentially getting statistics from the calibration set to “initialize” the quantization parameters. After this, you can fine-tune on the full data, or even a small subset if the model is well adapted to quantization.
• Representative datasets help in better generalization and reduce over-adaptation.
• If the dataset is small-medium, QAT can be performed during training directly.

Understanding Post-Training Quantization (PTQ)

PTQ is a conversion technique applied to a pre-trained model to reduce model size and improve inference speed without requiring retraining.

It lowers model weight and activation precisions to 8 bits (INT8), which for an FP16 model is just a half of the original footprint, leading to a significant improvement in inference speed.

PTQ algorithms often determine the quantization scaling factors with limited calibration data through a simple parameter space search without training or fine-tuning. Metrics for searching scaling factors include MSE distance and cosine distance.

The general process of PTQ involves the following steps:

1. Calibration:

• Representative Dataset: PTQ relies on a representative dataset (also called a calibration dataset) to collect statistics about the activations of the trained floating-point model.
• Forward Passes: The calibration dataset is passed through the (unquantized) model, and Observer modules (similar to those used in QAT) are inserted at various points to collect statistics. These observers typically track the minimum, maximum, and potentially other distribution information (e.g., histograms, percentiles) of activations.
• No Weight Updates: Importantly, no training or weight updates occur during this calibration phase. It’s purely about observing the model’s behavior on representative data.

2. Quantization Parameter Calculation:

• Statistics Analysis: The statistics collected by the observers are analyzed to determine the optimal quantization parameters (scale and zero_point) for each layer or tensor. The goal is to choose quantization parameters that minimize the difference between the floating-point and quantized representations of the activations and weights.
• Calibration Methods: Various calibration methods can be used such as min-max, using entropy.
• Percentile: Using percentiles (e.g. 99, 99.9th) of the observed values to determine the quantization range, making the quantization more robust to outliers.

3. Model Conversion:

• Quantizing Weights: The weights of the model are quantized using the calculated scale and zero_point values. This usually involves rounding the weights to the nearest quantized integer value.
  Note, in the context of LLM inference, memory is one of the major bottlenecks, therefore weight-only quantization is popular as weights make up the majority of model’s memory usage. If compute is the bottleneck for inference, weight only quantization can add more overhead due to dequantization step.
• Quantizing Activations: The activation functions might be modified or replaced to produce quantized outputs. This might involve adding quantization operations (e.g., quantize_per_tensor, quantize_per_channel) after activation functions.
  Note, activations can guide which weights should be quantized. Activation-Aware Quantization (AWQ) introduced in the paper https://arxiv.org/pdf/2310.09259. They key idea behind AWQ is that not all weights contribute equally to the model’s performance. Some weights are more “salient” and need higher precision. AWQ identifies these salient weights (0.1–1%) by observing activation magnitudes and keeps them in higher precision.
• Inserting Quantization/Dequantization Operations (Optional): Depending on the target hardware and framework, explicit quantization and dequantization operations might be inserted into the model graph to handle the conversion between integer and floating-point representations during inference. Some frameworks can automatically handle this.

When to Choose QAT over PTQ

• When the model architecture is sensitive to quantization (activations) and retraining is feasible.
• Accuracy sensitive: When high accuracy is paramount, regardless of higher training costs.
• Low bit quantization: Low bit quantization like 2/4 bit.

When to Choose PTQ over QAT

• Large model size: When retraining is impractical or computational resources are limited.
• Accuracy sensitivity: When a slight reduction in accuracy is acceptable.
• A well-tuned PTQ model can be a good base for subsequent QAT fine-tuning.

Best of both worlds (PTQ with QAT Fine-tuning):

Combining PTQ and QAT to get best of both worlds is an area of active research.

Step 1: PTQ Initialization:

• Start with a pre-trained, full-precision model.
• Apply PTQ to quantize the model (weights, or both weights and activations) to the desired lower precision (e.g. 8-bit, 4-bit).
• This quantized model serves as the starting point (initialization) for the QAT fine-tuning stage.

Step 2: QAT Fine-tuning:

• Fine-tune the PTQ-quantized model using QAT.
• During fine-tuning, the model’s weights and activations are simulated in their quantized form (using techniques like “fake quantization”), allowing the model to learn to compensate for the quantization errors.
• This fine-tuning step is typically much shorter than full QAT from scratch, as the model is already partially adapted to quantization thanks to the PTQ initialization.
• Update the weights using backpropagation. Since the model is already in the quantized space, only a few epochs of fine-tuning are sufficient to recover the loss in accuracy.

PTQ with QAT Fine-tuning leverages strengths of both the techniques — improved accuracy and is faster and cheaper than full QAT.

QAT and PTQ in Recommender Systems (RecSys)

RecSys models often involve large embedding tables to represent categorical features, making them memory-intensive.

PTQ can be a good starting point for quantizing embedding tables, especially for 8-bit quantization or when resources are limited. However, for lower bit precision (2/4 bit) or when accuracy is paramount, QAT or more sophisticated quantization techniques are generally preferred.

Sensitivity analysis is important to determine if PTQ is suitable for embedding tables.

QAT and PTQ in Large Language Models (LLMs)

LLMs, with their billions of parameters, benefit significantly from quantization for deployment on various devices. Quantization enables LLMs to run on a wider range of devices, including single GPUs or even CPUs, by reducing their memory and storage requirements.

• PTQ is commonly used to quantize LLMs after training due to their large size, reducing their size and improving inference speed.
• QAT can be applied to LLMs to fine-tune them with quantization awareness, leading to better accuracy preservation compared to PTQ.

As we see a lot boils down to the sensitivity of layers, computational budget and application goals. Here is a guide for general sensitivity of different layers for quantization.

Numeric sensitivity analysis is important to understand the impact of quantization per layer or groups of layer. Pytorch Numeric suite is a toolkit to do that.

Recsys layers:

Press enter or click to view image in full size

LLM layers:

Press enter or click to view image in full size

Conclusion

The choice between QAT and PTQ depends on the specific requirements of the application (quantization type, deployment environment, accuracy sensitivity) and the available resources. QAT generally offers better accuracy but requires more training resources. PTQ is more efficient in terms of training time and resources but may result in some accuracy loss.