MobileQuant: Mobile-friendly Quantization for On-device Language Models

Previous research in post-training quantization for LLMs can be categorized into three main groups:

Quantization aware training (QAT) involves retraining or fine-tuning full-precision models using differentiable quantizers. Recent research Liu et al. (2023); Bondarenko et al. (2023) has shown that QAT outperforms PTQ methods, particularly with in-domain training data. However, QAT requires extensive training, which is often impractical for LLMs. Additionally, QAT may be vulnerable to domain shifts if the data used for pretraining is unavailable. In contrast, our approach is zero-shot, only requiring a minimal set of calibration samples and a limited compute budget. Once trained, our model remains mathematically equivalent to the original model when unquantized, enhancing its adaptability to various downstream tasks.

Quantization methods are differentiated by several main design choices, with varying levels of hardware support. In this section, we first list these design choices and then highlight the limitations of existing works with respect to these choices.

State-of-the-art quantization methods demonstrate strong performance on a server use case (i.e. high-end GPU). However, they either utilize on-the-fly dequantization and 16-bit floating point operations Frantar et al. (2023); Lin et al. (2024), which are computationally inefficient, or dynamic per-token quantization Xiao et al. (2023); Shao et al. (2024), which, as previously mentioned, has no support on edge devices.

We, instead, consider design choices that are widely supported and optimized on modern edge devices (e.g. Mobile NPUs), namely i) fixed-point weight and activation quantization with integer arithmetic operations, and ii) per-tensor/channel quantization with static pre-computed ranges. Our objective is hence to improve existing state-of-the-art approaches such as SmoothQuant Xiao et al. (2023) and OmniQuant Shao et al. (2024) while staying within the limits of hardware support on device.

Prior efforts on LLM quantization Dettmers et al. (2022); Xiao et al. (2023) observed that activations are harder to quantize compared to the model weights due to the outlier channel dimensions with diverse min-max ranges. As an example, given a fully connected layer $\mathbf{Y}=\mathbf{XW}$, $\mathbf{W}\in\mathbb{R}^{N\times M},\mathbf{X}\in\mathbb{R}^{N},\mathbf{Y}\in\mathbb{R}^{M}$, specific channel dimensions $\{i:0\leq i<N\}$ in $\mathbf{X}$ may have a wide min-max range across different data samples, causing substantial quantization errors. To counteract this, previous methods proposed a weight equivalent transform defined by a scaling vector $\mathbf{S}\in\mathbb{R}^{N}$:

The goal is to find the optimal scaling vector $\mathbf{S}$ such that both $\hat{\mathbf{X}}$ and $\hat{\mathbf{W}}$ are easier to quantize compared to the original $\mathbf{X}$ and $\mathbf{W}$. SmoothQuant Xiao et al. (2023) reparameterized $\mathbf{S}$ as $\mathbf{s}_{i}=\frac{max(|\mathbf{X}_{i}|)^{\alpha}}{max(|\mathbf{W}_{i}|)^{(1-\alpha)}})$, $0\leq i<N$, and searched for the hyper-parameter $\alpha$. The obtained $\mathbf{S}$ is similar to the closed-form solution derived in Nagel et al. (2019). OmniQuant Shao et al. (2024) extended SmoothQuant Xiao et al. (2023) by learning $\mathbf{S}$, together with the weight clipping parameters via block-wise error minimization. Here, both $\mathbf{S}$ and $\mathbf{S}^{-1}$ can be fused to the adjacent linear layers, making the transformation mathematically equivalent to the original models. Figure 1 provides an illustration of the transformation among consecutive linear layers.

The weight equivalent transformation approaches used in SmoothQuant and OmniQuant, as described in Section 3.2, demonstrate strong performance on GPU-like hardware. However, they do not work out of the box for edge devices. Specifically, two challenges remain: i) the weight transformations cannot propagate beyond non-linear operators, e.g. Softmax, RMSNorm Zhang and Sennrich (2019), LayerNorm Ba et al. (2016), SiLU/GELU Hendrycks and Gimpel (2016). To counteract this, we apply weight transformations on all consecutive layers with linear components, e.g. between linear layers or affine transformations in the normalization layers, while keeping the nonlinear activations in 16-bit integers; ii) with the weight transformation, the distribution of the activations shifts accordingly. This causes essential difficulty for learning-based approaches like OmniQuant Shao et al. (2024), when the min-max range for the activations changes after each training iteration. OmniQuant Shao et al. (2024) proposed to bypass the issue with dynamic per-token quantization, which has no hardware support on-device.

Given the distribution of the activations shifts accordingly with the weight transformation, the ideal solution is to re-estimate the activation ranges across the training set after each training iteration. However, doing so is computationally prohibited. Hence, we propose to learn the activation range jointly with the weight transformation. Given an activation tensor $\mathbf{X}$, instead of learning the min and max values $f_{min}(\mathbf{X})$, $f_{max}(\mathbf{X})$ directly, we leverage the correlation between $f_{min}$, $f_{max}$ and the scale and offset parameters, $\alpha,\beta\in\mathbb{R}$, for quantization. With the targeted bit-width $bw$, quantizing $\mathbf{X}$ can be formulated as:

Here, $\mathbf{X}_{int}$ refers to the quantized tensor of $\mathbf{X}$, $ste$ refers to straight-through estimator. We can therefore learn ${\color[rgb]{0.4375,0,0}f_{min}}={\color[rgb]{0,0.88,0}\alpha}{\color[rgb]{0,0.88,0}\beta}$ and ${\color[rgb]{0.4375,0,0}f_{max}}={\color[rgb]{0,0.88,0}\alpha}q_{max}+{\color[rgb]{0,0.88,0}\alpha}{\color[rgb]{0,0.88,0}\beta}$ indirectly by learning ${\color[rgb]{0,0.88,0}\alpha}$ and ${\color[rgb]{0,0.88,0}\beta}$, which are computationally more stable.

To learn the equivalent transformation, previous works either resort to closed-form solutions Nagel et al. (2019), search-based methods Xiao et al. (2023); Lin et al. (2024), or layer-wise error minimization Shao et al. (2024). These solutions require limited training budget, but, as shown in Section. 5.4, lead to sub-optimal performance. Particularly, given the restricted form of supervision, we show that these methods cannot scale with more training samples or iterations. We, instead, propose to jointly optimize all the training parameters, including the weight equalization parameters $\mathbf{S}$, weight clipping parameters used in OmniQuant Shao et al. (2024), and the range parameters $\alpha,\beta$ for all layers in an end-to-end manner. Compared to previous PTQ approaches, which struggle with more training samples and epochs, we demonstrate that our holistic optimization approach consistently improves the performance with larger training settings for different LLM architectures. Compared to QAT, our method preserves model generalizability and does not overfit to specific calibration samples, achieving near-lossless zero-shot performance.

We perform experiments by training and simulating the quantization on GPUs and further evaluate the on-device performance on a Samsung Galaxy S24, with the Snapdragon 8 Gen 3 HTP as the compute unit. All models were trained on two A100 GPUs, with a maximum sequence length of 2048.

In this section, we extend state-of-the-art weight-activation quantization methods, SmoothQuant Xiao et al. (2023) and OmniQuant Shao et al. (2024) on device and use them as baselines. As these approaches utilize dynamic per-token quantization for the activation, which is not supported on edge devices, we modify these methods to work on device by using static per-tensor activation quantization, referring to these variants as OmniQuant-Static and SmoothQuant-Static respectively. Note that, for SmoothQuant, we only include evaluations on W8A8, which is the default setting used in the original work.

As shown in Table 1, both “Static” variants suffer from large performance degradation when evaluated on WikiText Merity et al. (2016). We further observe that the performance drop is mainly caused by quantizing the activations for the last linear layer in each MLP head (i.e. $\mathit{down\_proj}$ in LLaMA-like Touvron et al. (2023) models). To further alleviate this issue, we introduce an extra weight equalization transformation between consecutive linear layers in each MLP head (i.e. $\mathbf{S}$ between the $\mathit{up\_proj}$ and $\mathit{down\_proj}$ layers in TinyLLaMA Zhang et al. (2024)). The new models, which we termed SmoothQuant-Edge and OmniQuant-Edge respectively, significantly alleviate the performance degradation. For the remainder of this section, we use these adapted models as strong on-device baselines.

Table 1 shows that the learning-based approach, OmniQuant Shao et al. (2024), outperforms the search-based method, SmoothQuant Xiao et al. (2023), for all models by a notable margin. However, learning to transform the weights with fixed activation ranges is suboptimal, as the activation ranges shift after each training iteration. We further evaluate the impact of incorporating activation range learning (ARL), described in Section. 4.2, into OmniQuant Shao et al. (2024). In other words, we learn the per-tensor scale and offset parameters, together with the weight transformation via block-wise error minimization.

Table 3 demonstrates that activation range learning (ARL) consistently improves the performance for all LLM models across all settings. The gains are larger for quantized models exhibiting a larger performance gap compared to the FP16 models. Notably, these models require more training steps to mitigate the quantization errors, leading to larger range shifts for the activation.

In the previous section, we show that incorporating ARL into our baselines results in consistent improvements. Nonetheless, there is still a notable performance gap between the quantized models and the FP16 models, especially under the W4A8 setting. In order to reduce this gap, we attempt to improve the performance by scaling up the performance, namely increasing the number of calibration samples and the number of training epochs. However, Table 2 shows that the performance of all considered models saturate as we scale the training up using the block-wise approach proposed in OmniQuant Shao et al. (2024). We therefore conjecture that the optimization is hindered by the block-wise error minimization objective that provides limited global supervision. To verify this, we use our end-to-end training pipeline and jointly optimize all trainable parameters of the whole model, namely the weight transformation, clipping, and activation range learning parameters.

As shown in Table 2, our end-to-end trained models demonstrate consistent improvements with more training samples and iterations, only underperforming the blockwise optimized models in the smallest settings when the models were undertrained. We currently train the models with up to 1024 samples for 60 epochs but posit that the models could be further improved with more diverse samples and larger training settings.

Following previous approaches Xiao et al. (2023); Shao et al. (2024); Liu et al. (2023), we perform zero-shot evaluations on representative tasks from the Harness benchmark Gao et al. (2023). Table 4 shows that, in addition to the WikiText perplexity, our method also improves the quantization performance for the common sense reasoning tasks in general, without using any in-domain data. The improvements are consistent for most benchmarks and we believe that the performance of our method could be further improved with in-domain data, especially for benchmarks with a large domain shift relative to our calibration set (i.e. Pile Gao et al. (2020)).