EfficientQAT: Efficient Quantization-Aware Training for Large Language Models

In this section, we introduce EfficientQAT, a novel quantization-aware training framework for LLMs that enhances memory efficiency. As illustrated in Figure 3, traditional QAT approaches train the full-precision weights $\mathbf{W}$ and quantization parameters $s$ (step sizes) and $z$ (zero points) simultaneously in an end-to-end manner, which significantly increases the memory requirements due to the large number of parameters involved. To address this issue, EfficientQAT adopts a two-stage strategy: block-wise training of all parameters (Block-AP) and end-to-end training of quantization parameters (E2E-QP). In the Block-AP phase, model parameters and quantization parameters are trained block-by-block using reconstruction loss, which not only allows for precise calibration with full training but also reduces memory consumption brecq ; omniquant by block-wise training. Following this, the E2E-QP phase fixes the quantized weights and trains the step sizes exclusively on target datasets, thus achieving quantization-aware training in a memory-efficient way. Details on Block-AP and E2E-QP are further described in Sections 3.2 and 3.3, respectively.

In this section, we introduce the Block-Wise Training of All Parameters (Block-AP) approach, designed to efficiently provide an effective initialization for following end-to-end training.

Quantization and Dequantization. Specifically, Block-AP begins with a standard uniform quantization method:

where $\lfloor\cdot\rceil$ represents the rounding operation. $N$ is the target bit number. $\mathbf{W}_{int}$ and $\mathbf{W}$ denote the quantized integer and full-precision weights, respectively. $s$ is the scaling factor and $z$ is the zero point. In the forward propagation, the quantized weights are converted back to full precision as follows:

Here, $\widehat{\mathbf{W}}$ refers to the dequantized weights used in the forward computation. The processes of quantization (Eq.(1)) and dequantization (Eq.(2)) are integrated within the computation graph and can be optimized through gradient descent in a quantization-aware manner.

Blcok-wise Quantization-aware Training. Traditional QAT methods bitnet158 ; lsq ; llmqat train the entire network using Eq.(1) and Eq.(2) in an end-to-end fashion, which typically requires substantial computational resources and extensive data to prevent overfitting. Here we aim to enhance the training efficiency of QAT. Previous studies, such as BRECQ brecq , have demonstrated that block-wise training achieves faster convergence and requires less training time, data, and memory than end-to-end training given a pre-trained model. Following the methodologies in BRECQ brecq and OmniQuant omniquant , Block-AP sequentially conducts quantization-aware training within one transformer block before moving on to the next under a block-wise reconstruction framework.

Full Training of Model Weights and Quantization Parameters. Unlike previous methods which optimize several quantization parameters such as rounding parameters adaround ; weight_round ; flexround , clipping parameters omniquant , and step sizes lsq ; cbq , Block-AP behaves like QAT, training all inherent parameters from Eq.(1) and Eq.(2), including scaling factor $s$, zero point $z$, and model weights $\mathbf{W}$.

In our Block-AP approach, a straightforward full-training regimen outperforms existing partial-training variants adaround ; brecq ; cbq with intricate designs. Traditional training methods involving rounding parameters adaround ; brecq ; cbq serve as regularization techniques, constraining the update range of integral weights to $(-1,+1)$ to mitigate overfitting. However, this approach limits the solution space, potentially hindering the final performance of quantized models. By simply scaling the training samples from $128$ to $4096$, Block-AP effectively addresses the overfitting issue (see Figure 5 for details) while improving performance compared to training with rounding parameters. Our empirical findings demonstrate the superiority of full training within our Block-AP over existing partial-training variants adaround ; brecq ; cbq , as shown in Table 7.

Following block-wise training, we obtain the quantized model which includes quantized weights $\mathbf{W}_{q}$, step sizes $s$, and zero points $z$ for each quantization group. The weights $\mathbf{W}_{q}$ and zero points $z$ are stored in a low-bit format, while step sizes $s$ are stored in FP16. Note that $s$ and $z$ are shared within their respective quantization groups and constitute only a small fraction of the model’s parameters, approximately 1.6% for a group size of 64. Moreover, the model’s memory footprint is substantially reduced by transitioning from full-precision 16-bit weights to 2/3/4-bit quantized weights.

We further introduce the End-to-End Training of Quantization Parameters (E2E-QP), aimed at efficiently training the entire quantized model on target datasets.

End-to-End Training of step sizes. Unlike traditional Quantization-Aware Training (QAT) methods llmqat ; bitnet158 that train full-precision weights, E2E-QP begins with $\mathbf{W}_{q}$ initialized via Block-AP and focuses solely on the training of quantization parameters ($s$ and $z$). Our findings indicate that training $s$, $z$, or both yields similar performance (see Table 7 for details). However, since training $z$ involves converting it from a low-bits format to full-precision, we typically train only $s$ by default unless specified otherwise to avoid additional memory overhead.

Additionally, within E2E-QP, there is no quantization process as per Equation (1); only the dequantization process occurs as described in Equation (2). Thus, the gradient of the trainable parameter $s$ is computed as $\frac{\partial\widehat{w}}{\partial s}=w_{q}-z$.

Overall, the memory usage for training in E2E-QP is drastically reduced due to the reduced trainable parameter count. Detailed memory footprints for various model sizes and bits under E2E-QP are listed in Table 8. For instance, the Llama-2-70B model can complete 2-bit QAT through E2E-QP using only 34.2GB of memory. Equipped with E2E-QP, EfficientQAT is adaptable to different scenarios by simply changing the training datasets, which includes applications such as continual pre-training and instruction-tuning alpaca .

Training. We conduct experiments on the Llama-2 and Llama-3 models. For Block-AP, we use 4096 samples from RedPajama redpajama with a context length of 2048. We train each block with batch size as 2 and epochs as 2, setting the learning rate of quantization parameters as $1\text{\times}{10}^{-4}$, and the learning rate of weights as $2\text{\times}{10}^{-5}$ for 2-bit and $1\text{\times}{10}^{-5}$ for 3/4-bits. For E2E-QP, we also employ 4096 samples from RedPajama redpajama but with a context length of 4096. We train the entire model with batch size as 32 and epoch as 1, and set the learning rate of step size as $2\text{\times}{10}^{-5}$ for 2-bit and $1\text{\times}{10}^{-5}$ for 3/4-bits.

Evaluation. We assess the zero-shot accuracy of five common-sense reasoning tasks using the v0.4.2 lm-evaluation-harness^†††https://github.com/EleutherAI/lm-evaluation-harness. The tasks include WinoGrande winogrande , PIQA piqa , HellaSwag hellaswag , Arc-Easy arc , and Arc-Challenge arc . We also measure the perplexity of Wikitext2 and C4 with a 2048 context length, as done in previous studies gptq ; omniquant .

PTQ Baseline. We compare our results with PTQ methods from uniform quantization such as GPTQ gptq , AWQ awq , OmniQ omniquant , and AutoRound weight_round , and vector quantization including QuIP# quip-sharp and AQLM aqlm .

PTQ Results. The accuracy and perplexity results are presented in Table 1 and Table 3, respectively. It seems that minimal performance degradation with 4-bit group quantization, and even simple round-to-nearest (RTN) quantization achieves respectable results. Therefore, EfficientQAT offers a slight perplexity improvement ($\sim$0.02) in 4-bit uniform quantization and comparable zero-shot accuracy with existing methods. As the number of quantization bits decreases, EfficientQAT’s performance advantage increases. In 3-bit quantization, EfficientQAT surpasses the optimization-based method AutoRound by approximately 0.5% in zero-shot accuracy and outperforms OmniQuant by 0.14 to 0.43 in perplexity. Notably, in 2-bit quantization, EfficientQAT outperforms AutoRound weight_round by 5% in zero-shot accuracy for Llama-2-7b. EfficientQAT also exceeds previous vector quantization methods like AQLM aqlm using a 2x8 codebook. While AQLM aqlm with a 1x16 codebook and QuIP# quip-sharp perform slightly better in 2-bit scenarios, they introduce notable computational overhead llm-qbench and can even reduce inference speed, as shown in Figure 2(a). Thus, EfficientQAT successfully pushes the limits of uniform quantization in LLMs, making hardware-friendly uniform quantization competitive with more complex vector quantization methods. However, we can also find that Llama-3 models suffer more performance degeneration after quantization llama3_quant , which may caused by the full training with 15T tokens other than 2T tokens.

QAT Baseline. We also compare our results with existing QAT methods, including LLM-QAT llmqat , BitDistiller bitdistiller , PB-LLM pbllm and DB-LLM dbllm .

QAT Results. As demonstrated in Table 2, EfficientQAT markedly surpasses other QAT methods in performance. Specifically, EfficientQAT reduces perplexity by 0.89 compared to BitDistiller bitdistiller when applied to the Llama-2-7B model with w2g128 quantization. In the more challenging scenario of Llama-3 quantization llama3_quant , EfficientQAT exceeds the performance of DB-LLM dbllm by 5.98 in C4 perplexity and by 8.57 points in average accuracy with the w2g128 quantization.

Training and Evaluation. Following existing works qa-lora ; ir-qlora , we train Llama-1 models on the Alpaca dataset alpaca and assess their performance by measuring average 5-shot MMLU mmlu accuracy works qa-lora ; ir-qlora . The training hyperparameters are identical to those described in Section 4.1, except we replace the RedPajama dataset redpajama with Alpaca. In line with QLoRA’s methodology qlora , we adjust the source context length to 384 and the target context length to 128, training for 10,000 steps with a batch size of 16.

Baseline. We benchmark EfficientQAT against several leading methods, including QLoRA qlora , QA-LoRA qa-lora , PEQA peqa , and IR-QLoRA ir-qlora , across quantization setting of 2, 3, and 4 bits. Consistent with QA-LoRA qa-lora , we also employ GPTQ gptq to quantize the fine-tuned QLoRA models into a low-bit format without FP16 LoRA for equitable comparison.

Results. Both Table 7 and Figure 1(b) indicate that EfficientQAT significantly outperforms existing Q-PEFT methods. For instance, in channel-wise quantization (group size of -1), EfficientQAT achieves more than 3% higher accuracy than PEQA peqa . In the 2-bit quantization scenario, the superiority of EfficientQAT is even more pronounced, surpassing QA-LoRA qa-lora by 5.1% and 4.0% in 7B and 13B models, respectively, and outperforming PEQA by 4.5% and 8.7% in the same models. Moreover, Table 7 also demonstrates that EfficientQAT outperforms both QA-LoRA and QLoRA with GPTQ in smaller model memory footprint (larger group size).

The EfficientQAT algorithm is comprised of two main components: Block-AP and E2E-QP. This section evaluates the effectiveness, trainable parameters, and training sample requirements of each component. We present the average perplexity for WikiText2 and C4 datasets, and the average accuracy for five zero-shot reasoning tasks, similar to Table 1.

Effectiveness of each component. As indicated in Table 7, both the Block-AP and E2E-QP components significantly enhance performance, with their combination yielding the best results. Notably, Block-AP outperforms E2E-QP, aligning with findings from BRECQ brecq .

Trainable parameters of Block-AP. Block-AP trains all parameters, including original weights and quantization parameters. Previous methods have introduced various training strategies to mitigate overfitting, such as trained rounding adaround ; weight_round , clipping thresholds omniquant , and step sizes lsq ; cbq . We compare Block-AP with these methods by modifying only the trainable parameters of Block-AP. As shown in Table 7, Block-AP (training $s$, $z$, $\mathbf{W}$) performs best with an acceptable training cost. Additionally, the memory footprint of directly training $\mathbf{W}$ is even smaller than that of training the rounding operation, which requires an additional copy of rounding parameters.

Trainable parameters of E2E-QP. We further examine the trainable parameters within E2E-QP. Table 7 shows that training $s$, $z$, or both yields similar performance. However, given that converting $z$ from an original low-bit representation to a trainable FP16 format increases the average bit count, we opt to train only $s$ by default.

Samples number of Block-AP. We assess the number of training samples for Block-AP, noting that E2E-QP trains all parameters, which may lead to overfitting. To address this, we introduce an additional 64 unseen samples from ReadPajama to evaluate the overfitting issue. We adjust the training epochs to ensure a similar total training time, allowing for fair comparisons across different sample sizes. As illustrated in Figure 5, increasing the number of training samples significantly reduces the gap between training loss and validation loss from 1.07 to 0.06. This reduction corresponds to an increase in the average accuracy for zero-shot tasks from 57.14% to 58.99%. Consequently, we set the default number of training samples for E2E-QP at 4096, as this maintains a minimal gap between training and validation losses.

Samples number of E2E-QP. In the E2E-QP, we train the model for 1 epoch to avoid over-fitting. Our examination of the training sample sizes for E2E-QP, detailed in Table 5, reveals that average perplexity consistently improves as sample sizes increase from 128 to 32,674. However, there is no significant improvement in average accuracy beyond 4096 samples. Therefore, we set the training sample size for E2E-QP at 4096 by default to balance efficiency and performance. Nonetheless, it is possible to further enhance the performance of EfficientQAT by increasing the sample size.

Training Efficiency Table 8 illustrates the required memory and time for training Lllama-2 models using EfficientQAT. The results indicate that the model completes training rapidly, taking 4.8 hours for the 7B model and 40.9 hours for the 70B model. Furthermore, the adoption of a quantized backbone in the end-to-end quantization process (E2E-QP) significantly reduces memory usage due to the lower bit width. Consequently, EfficientQAT accomplishes 2-bit training for the 70B model using only 34.2GB of memory.

Inference Efficiency Weight-only quantization reduces the memory requirements of LLMs and accelerates inference, as LLMs are predominantly memory-bound llm-unveil . Quantization modifies the computation in the linear layers while other operations remain unchanged. To evaluate this, we calculate the relative speedup in the linear layers’ processing. According to Table 9, INT2 quantization enhances the forward-pass speed by approximately 2.9x to 4.4x. Furthermore, due to the leverage of standard uniform quantization, the quantized models of EfficientQAT can also achieve speedup through other toolbox, such as MLC-LLM mlc-llm , AWQ awq , and BitBLAS bitblas .