FlattenQuant: Breaking Through the Inference Compute-bound for Large Language Models with Per-tensor Quantization

Yi Zhang,Fei Yang,Shuang Peng,Fangyu Wang,Aimin Pan

2024-02-28

Abstract:Large language models (LLMs) have demonstrated state-of-the-art performance across various tasks. However, the latency of inference and the large GPU memory consumption of LLMs restrict their deployment performance. Recently, there have been some efficient attempts to quantize LLMs, yet inference with large batch size or long sequence still has the issue of being compute-bound. Fine-grained quantization methods have showcased their proficiency in achieving low-bit quantization for LLMs, while requiring FP16 data type for linear layer computations, which is time-consuming when dealing with large batch size or long sequence. In this paper, we introduce a method called FlattenQuant, which significantly reduces the maximum value of the tensor by flattening the large channels in the tensor, to achieve low bit per-tensor quantization with minimal accuracy loss. Our experiments show that FlattenQuant can directly use 4 bits to achieve 48.29% of the linear layer calculation in LLMs, with the remaining layers using 8 bits. The 4-bit matrix multiplication introduced in the FlattenQuant method can effectively address the compute-bound caused by large matrix calculation. Our work achieves up to 2$\times$ speedup and 2.3$\times$ memory reduction for LLMs with negligible loss in accuracy.

Machine Learning,Artificial Intelligence,Computation and Language

What problem does this paper attempt to address?

The paper primarily addresses the computational bottleneck issues encountered by large language models (LLMs) during the inference process. Specifically, when handling large batches or long sequences, the inference process of large language models is often compute-bound, leading to a significant increase in inference latency. Although some quantization methods have been developed to optimize the inference performance of large language models, such as the GPTQ quantization method which alleviates memory bottlenecks through 4-bit quantized weights, they still face compute-bound issues when dealing with large batches or long sequences.

To tackle this challenge, the paper proposes a method called FlattenQuant. This method significantly reduces the maximum value of tensors by flattening high-value channels in the activation tensor and adding extra channels to accommodate these high values. In this way, even when processing large batches or long sequences, low-bit tensor quantization can be achieved while maintaining high accuracy. FlattenQuant enables the direct use of 4-bit quantization for 48.29% of the linear layer computations in large language models, with the remaining parts using 8-bit quantization. This method effectively addresses the compute-bound issue, achieving up to 2 times speedup and 2.3 times memory reduction compared to using FP16 data type, while ensuring minimal precision loss.

In summary, the problem the paper attempts to solve is: how to effectively reduce the computational bottlenecks faced by large language models when processing large batches or long sequences, in order to improve inference speed and reduce memory consumption. FlattenQuant provides a feasible solution for this.