FlatQuant: Flatness Matters for LLM Quantization
Authors: name: Yuxuan Sun; name: Ruikang Liu; name: Haoli Bai; name: Han Bao; name: Kang Zhao; name: Yuening Li; name: Jiaxin Hu; name: Xianzhi Yu; name: Lu Hou; name: Chun Yuan; name: Xin Jiang; name: Wulong Liu; name: Jun Yao
Published: 2024-10-12 08:10:28

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 1% accuracy drop for
W4A4 quantization on the LLaMA-3-70B model, surpassing SpinQuant by
7.5%. For inference latency, FlatQuant reduces the slowdown induced
by pre-quantization transformation from 0.26x of QuaRot to merely
0.07x, bringing up to 2.3x speedup for prefill and
1.7x speedup for decoding, respectively. Code is available at:
https://github.com/ruikangliu/FlatQuant.