AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration

Ji Lin; Jiaming Tang; Haotian Tang; Shang Yang; Wei-Ming Chen; Wei-Chen Wang; Guangxuan Xiao; Xingyu Dang; Chuang Gan; Song Han

Source-linked AI summary

AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration

Ji Lin, Jiaming Tang, Haotian Tang, Shang Yang, Wei-Ming Chen, Wei-Chen Wang, Guangxuan Xiao, Xingyu Dang, Chuang Gan, Song Han

arXiv:2306.00978v6Jun 1, 2023cs.CL

TL;DR

Large language models are difficult to deploy on edge devices because their size exceeds constrained hardware resources. AWQ uses activation-aware per-channel scaling for low-bit weight-only quantization, and TinyChat implements efficient 4-bit inference. The approach outperforms existing work across language, instruction-tuned, and multi-modal models while delivering 3.2-3.3× measured speedups over Huggingface FP16 implementations.

Problem
LLMs are challenging to deploy on edge devices because their large model sizes exceed limited hardware resources.
Method
AWQ uses activation statistics to identify salient channels and per-channel scaling to reduce quantization error without reconstruction, while TinyChat efficiently deploys 4-bit models.
Results
AWQ outperforms existing work across language-modeling tasks and model families, including instruction-tuned and multi-modal LMs, while TinyChat achieves 3.2-3.3× speedups over Huggingface FP16.
Takeaways & Limitations
AWQ preserves generalist abilities across domains and modalities while enabling practical low-bit LLM deployment on desktop and mobile edge hardware.
Takeaways & Limitations
Keeping salient weights in FP16 improves quantized performance but creates a hardware-inefficient mixed-precision implementation.

Abstract

from arXiv · show

Large language models (LLMs) have transformed numerous AI applications. On-device LLM is becoming increasingly important: running LLMs locally on edge devices can reduce the cloud computing cost and protect users' privacy. However, the astronomical model size and the limited hardware resource pose significant deployment challenges. We propose Activation-aware Weight Quantization (AWQ), a hardware-friendly approach for LLM low-bit weight-only quantization. AWQ finds that not all weights in an LLM are equally important. Protecting only 1% salient weights can greatly reduce quantization error. To identify salient weight channels, we should refer to the activation distribution, not weights. To avoid the hardware-inefficient mix-precision quantization, we mathematically derive that scaling up the salient channels can reduce the quantization error. AWQ employs an equivalent transformation to scale the salient weight channels to protect them. The scale is determined by collecting the activation statistics offline. AWQ does not rely on any backpropagation or reconstruction, so it generalizes to different domains and modalities without overfitting the calibration set. AWQ outperforms existing work on various language modeling and domain-specific benchmarks (coding and math). Thanks to better generalization, it achieves excellent quantization performance for instruction-tuned LMs and, for the first time, multi-modal LMs. Alongside AWQ, we implement TinyChat, an efficient and flexible inference framework tailored for 4-bit on-device LLM/VLMs. With kernel fusion and platform-aware weight packing, TinyChat offers more than 3x speedup over the Huggingface FP16 implementation on both desktop and mobile GPUs. It also democratizes the deployment of the 70B Llama-2 model on mobile GPUs.

1 INTRODUCTION

On-device LLM deployment promises lower latency, offline operation, reduced cloud costs, and improved data security, but model size makes deployment difficult. AWQ addresses this with activation-aware weight quantization, while TinyChat turns 4-bit compression into practical edge inference speedups.

Motivation: On-device LLMs reduce cloud delays and costs while supporting offline, real-time applications and keeping sensitive data local.These benefits are especially relevant to virtual assistants, chatbots, and autonomous vehicles.
Motivation: GPT-3's 175B parameters occupy 350GB in FP16, exceeding the 192GB memory of an NVIDIA B200 GPU.Edge devices have even more constrained memory resources.
AWQ: AWQ identifies salient weight channels from activation distributions and scales them per channel to reduce quantization error without mixed precision.The method protects a small fraction of important weights while retaining hardware-friendly full-weight quantization.
TinyChat: TinyChat combines on-the-fly dequantization, 4-bit weight packing, and kernel fusion to reduce inference overhead on byte-aligned hardware.The framework targets intermediate DRAM access and kernel launch overhead in linear layers.
Results: AWQ outperformed existing methods across tasks, model families, and sizes, including instruction-tuned and multi-modal language models.The reported evaluations include LLaMA, OPT, Vicuna, and OpenFlamingo.
Results: 3.2-3.3× average speedup over Huggingface FP16 was observed across desktop, laptop, and mobile GPUs.The system also enabled Llama-2-70B deployment on a single 64GB NVIDIA Jetson Orin.

2 RELATED WORK

LLM quantization research includes both activation-and-weight INT8 quantization and low-bit weight-only quantization, with this work focusing on the latter. Existing systems and methods address memory or inference costs, but hardware support and calibration generalization remain concerns.

Quantization methods: Quantization-aware training uses backpropagation, whereas post-training quantization operates after training and is commonly used for LLMs.PTQ is typically preferred because QAT incurs high training cost.
LLM quantization: W8A8 quantizes activations and weights to INT8, while W4A16 quantizes only weights into low-bit integers.The paper focuses on weight-only quantization because it reduces memory barriers and speeds token generation for memory-bound workloads.
LLM quantization: GPTQ uses reconstruction and second-order information but may overfit calibration data and weaken generalist abilities across domains and modalities.GPTQ also requires a reordering trick for some models, including LLaMA-7B and OPT-66B.
System support: Existing systems provide INT3 or INT4 kernels, group-wise INT4 quantization, or FP16×INT4 GEMM to reduce inference and I/O costs.The cited systems include GPTQ, GPTQ-for-LLaMA, FlexGen, llama.cpp, exllama, and FasterTransformer.

3 AWQ: ACTIVATION-AWARE WEIGHT QUANTIZATION

AWQ protects salient weight channels identified from activation distributions, then replaces hardware-inefficient mixed precision with equivalent per-channel scaling. Its search-based scaling reduces quantization error while preserving hardware efficiency and generalization.

3.1 Improving LLM Quantization by Preserving 1% Salient Weights: AWQ observes that a small fraction of salient weight channels disproportionately affects LLM performance, while selecting channels by weight norm is ineffective.Keeping 0.1%-1% of important channels in FP16 improves quantized performance, but activation distributions identify those channels more effectively than weight distributions.
3.1 Improving LLM Quantization by Preserving 1% Salient Weights: Mixed-precision protection improves quantized performance but makes system implementation difficult, motivating a hardware-friendly alternative.The limitation arises even when retaining only 0.1% of weights in FP16 adds little total model size.
3.2 Protecting Salient Weights by Activation-aware Scaling: AWQ reduces salient-channel quantization error through per-channel scaling rather than retaining those weights in FP16.Scaling preserves an equivalent computation by scaling weights and inversely scaling activations, avoiding the hardware inefficiency of mixed precision.
3.2 Protecting Salient Weights by Activation-aware Scaling: For OPT-6.7B, scaling 1% salient channels from s = 1 to s = 2 improves perplexity from 23.54 to 11.92.Larger scales reduce salient-channel relative error but can amplify non-salient-channel error, so the best perplexity occurs at s = 2.
3.2 Protecting Salient Weights by Activation-aware Scaling: AWQ searches per-input-channel scales using cached calibration activations, with α balancing salient- and non-salient-channel protection.The objective minimizes post-quantization output difference, and α is selected by a grid search over [0, 1].
3 AWQ: ACTIVATION-AWARE WEIGHT QUANTIZATION: On-device generation is memory bound because weight access dominates memory traffic, making weight-only quantization a suitable deployment target.The supplied bottleneck analysis reports that W4A16 quantization can improve arithmetic intensity by 4×.
3 AWQ: ACTIVATION-AWARE WEIGHT QUANTIZATION: AWQ requires neither regression nor backpropagation and uses only average per-channel activation magnitudes, reducing calibration-set overfitting.The method is described as preserving knowledge outside the calibration distribution while requiring fewer quantization data.

4 TINYCHAT: MAPPING AWQ ONTO EDGE PLATFORMS

TinyChat converts AWQ’s theoretical 4-bit memory savings into practical edge-device acceleration through fused kernels, on-the-fly dequantization, and platform-specific weight packing. It targets the memory-bound generation phase and achieves over 3× speedup across GPU platforms.

Deploy AWQ with TinyChat: SIMD-aware packing reorders weights according to device bit width, enabling runtime byte unpacking with bitwise operations on ARM NEON’s 128-bit registers.
Why AWQ Helps Accelerate On-Device LLMs: Generation is substantially slower than context processing, taking 310 ms for 20 generated tokens versus 10 ms to summarize a 200-token prompt.
Why AWQ Helps Accelerate On-Device LLMs: The FP16 generation workload is memory-bound because its arithmetic intensity is approximately 1, while weight access dominates memory traffic.
Why AWQ Helps Accelerate On-Device LLMs: 4-bit weight quantization raises the theoretical peak performance by 4× by reducing weight precision and increasing arithmetic intensity.
Deploy AWQ with TinyChat: More than 3× speedup is achieved over Huggingface FP16 across different LLM families on GPUs using 4-bit AWQ and TinyChat.
Deploy AWQ with TinyChat: TinyChat fuses dequantization with matrix multiplication to avoid writing dequantized weights to DRAM for both matrix-matrix and matrix-vector products.
Deploy AWQ with TinyChat: Kernel fusion combines operations in layer normalization and attention, including QKV projections, positional embeddings, and KV-cache updates.

5 EXPERIMENTS

Experiments evaluate AWQ across model families, instruction-tuned and multimodal models, complex-generation tasks, calibration settings, and TinyChat deployment. AWQ generally preserves or improves quantized quality, while TinyChat delivers substantial inference speedups and broader edge deployment.

Language modeling: AWQ consistently outperforms RTN and GPTQ across LLaMA model scales from 7B to 70B, while also improving INT2 performance when combined with GPTQ.The paper reports AWQ as orthogonal to GPTQ, enabling further gains under extreme low-bit quantization.
Instruction-tuned models: AWQ improves instruction-tuned Vicuna models over RTN and GPTQ under INT3-g128 for both 7B and 13B models.The comparison uses GPT-4 evaluation with 160 trials.
Multimodal models: AWQ reduces OpenFlamingo-9B 32-shot quantization degradation from 4.57 to 1.17 under INT4-g128 while providing 4× model-size reduction.It also achieves lossless performance for VILA across 11 visual-language benchmarks and improves qualitative LLaVA reasoning responses over RTN.
Programming and mathematics: AWQ outperforms existing methods on MBPP and GSM8K, with INT4-g128 performance comparable to the original FP16 models.These results cover programming and multistep mathematical reasoning.
Inference deployment: TinyChat achieves 2.7-3.9× speedup on RTX 4090 and about 3× speedup for VILA-7B and VILA-13B on Jetson Orin.It also runs Llama-2-13B at 33 tokens/s on an 8GB RTX 4070, where the FP16 implementation cannot fit 7B models, and reaches up to 1.7× over llama.cpp on Orin.

6 CONCLUSION

The paper concludes that AWQ provides effective low-bit weight-only compression by scaling salient channels, while preserving generalist abilities across domains and modalities. TinyChat converts AWQ’s memory savings into measured edge-inference speedups.

6 CONCLUSION: AWQ uses per-channel scaling to reduce quantization loss for salient weights without overfitting the calibration set.The method is presented as simple, effective, and applicable to low-bit weight-only LLM compression.
6 CONCLUSION: AWQ outperforms existing methods on language modeling and extends to instruction-tuned and multimodal language models.The conclusion attributes this scope to preserving generalist abilities across domains and modalities.
6 CONCLUSION: 3.2-3.3× measured speedups over Hugging Face FP16 implementations are achieved on desktop and mobile GPUs with TinyChat.The system translates AWQ’s theoretical memory savings into practical edge deployment benefits.