FlatQuant: Flatness Matters for LLM Quantization

Yuxuan Sun,Ruikang Liu,Haoli Bai,Han Bao,Kang Zhao,Yuening Li,Jiaxin Hu,Xianzhi Yu,Lu Hou,Chun Yuan,Xin Jiang,Wulong Liu,Jun Yao

2024-10-12

Abstract:Recently, quantization has been widely used for the compression and acceleration of large language models~(LLMs). Due to the outliers in LLMs, it is crucial to flatten weights and activations to minimize quantization error with the equally spaced quantization points. Prior research explores various pre-quantization transformations to suppress outliers, such as per-channel scaling and Hadamard transformation. However, we observe that these transformed weights and activations can still remain steep and outspread. In this paper, we propose FlatQuant (Fast and Learnable Affine Transformation), a new post-training quantization approach to enhance flatness of weights and activations. Our approach identifies optimal affine transformations tailored to each linear layer, calibrated in hours via a lightweight objective. To reduce runtime overhead, we apply Kronecker decomposition to the transformation matrices, and fuse all operations in FlatQuant into a single kernel. Extensive experiments show that FlatQuant sets up a new state-of-the-art quantization benchmark. For instance, it achieves less than $\textbf{1}\%$ accuracy drop for W4A4 quantization on the LLaMA-3-70B model, surpassing SpinQuant by $\textbf{7.5}\%$. For inference latency, FlatQuant reduces the slowdown induced by pre-quantization transformation from 0.26x of QuaRot to merely $\textbf{0.07x}$, bringing up to $\textbf{2.3x}$ speedup for prefill and $\textbf{1.7x}$ speedup for decoding, respectively. Code is available at: \url{<a class="link-external link-https" href="https://github.com/ruikangliu/FlatQuant" rel="external noopener nofollow">this https URL</a>}.

Computation and Language,Machine Learning

What problem does this paper attempt to address?

This paper attempts to solve the problem of non - flat weight and activation distributions encountered during the quantization process of large - language models (LLMs). Specifically, the paper points out that due to the large number of outliers in LLMs, this leads to an increase in quantization error. Although traditional pre - quantization conversion methods (such as channel scaling and Hadamard transformation) can reduce the influence of these outliers to a certain extent, they still cannot completely solve the problems of steepness and dispersion in weight and activation distributions. These problems not only affect the quantization accuracy but also increase the propagation of errors between Transformer layers.

To solve the above problems, the paper proposes a new post - training quantization method - FLATQUANT (Fast and Learnable Affine Transformation). FLATQUANT enhances the flatness of weights and activations by identifying the optimal affine transformation for each linear layer. This method can not only reduce quantization error more effectively but also significantly improve the performance of the quantized model in terms of inference speed. The paper verifies the effectiveness of FLATQUANT through extensive experiments and shows its superior performance in multiple benchmark tests. In particular, in low - bit quantization settings, FLATQUANT can achieve an accuracy drop of less than 1% while bringing a 2.3 - fold and 1.7 - fold speedup in the pre - filling and decoding stages respectively.