AWQ: Activation-aware Weight Quantization for on-device LLM compression and acceleration

https://arxiv.org/pdf/2306.00978

https://github.com/mit-han-lab/llm-awq

GitHub - mit-han-lab/llm-awq: [MLSys 2024 Best Paper Award] AWQ: Activation-aware Weight Quantization for LLM Compression and Ac

 

MLSys 2024 Best Paper Award

 

 

Abstract

Activation-aware Weight Quantization (AWQ): 중요한 가중치만 잘 살리면 LLM도 가볍게 돌아간다.

 

핵심 아이디어는 '모든 weight가 중요한 게 아니다' 입니다.

상위 1% 정도의 weight만 잘 보호하면, 전체 quantization error를 크게 줄일 수 있다는 사실을 발견했습니다.

 

그런데 어떤 weight가 중요한지를 판단하려면 weight 값 자체보다 해당 weight가 처리하는 activation의 분포를 보는 게 훨씬 더 효과적이라는 것이 저자들의 주장입니다.

 

일반적으로 사용되는 mixed-precision 방식은 하드웨어 구현에 비효율적인데, AWQ는 이를 피하면서도 중요한 weight 채널만 스케일링을 통해 효과적으로 보호합니다. 

 

또한, AWQ가 backpropagation이나 reconstruction을 전혀 사용하지 않습니다. 대신 offline으로 측정한 activation 통계만을 활용하므로, 다양한 도메인에 잘 일반화되고, calibration set에 과적합되지 않는다는 장점이 있습니다.

 

추가적으로 저자들은 TinyChat이라는 추론 프레임워크를 제안했습니다.

이 추론 프레임워크는 4bit LLM이 이론적으로 절감하는 메모리 사용량을 실제 속도 향상으로 연결시켜주는 프레임워크입니다.

• On-the-fly dequantization
• efficient 4-bit weight packing
• kernel fusion

그 결과, 가중치를 4비트로 줄였을 때 발생할 수 있는 성능 손실 없이, 속도 향상을 달성했습니다.

 

Method

(a) 기존 RTN(Round To Nearest) quantization, (b) 1%의 중요한 weight를 보존하고 quantization, (c) quantizatione rror를 줄이기 위한 최적의 scaling을 찾는 방법

 

기존 방법 RTN (a)는 PPL이 엄청나게 올라가는 문제가 있습니다.

 

이를 해결하기 위해, 중요한 weight를 보존하여 quantization (b) 한다면, PPL은 줄어들 수 있지만, mixed-precision 문제로 인해서 하드웨어 최적화가 어렵습니다.

 

이에 저자들은 quantization하기전에 weight를 scaling하여 해결할 수 있는 방법을 제안합니다.

 

(b) Improving LLM Quantization by Preserving 1% Salient Weights

저자들은 LLM의 weight가 모두 동일하게 중요하지 않다는 것을 관찰했습니다.

* 이는 매우 작은 비율의 salient weights가 다른 weight들에 비해서 LLM 성능에 훨씬 더 중요하다는 의미입니다.

 

Quantization할때, 이러한 중요한 weight는 건너뛰는 것이 추가적인 학습이나 regression없이 performance 손실을 최소화할 수 있는 방법입니다.

 

위 Table 1에서 볼 수 있듯이, FP16에서 0.1% ~ 3%의 비율을 보존하고 quantization하는게 기존 RTN 방법보다 훨씬 낮은 perplexity를 보였습니다. 또한 보존하는 weight를 선택할때 Activation, Weight(L2 norm), Random에 대해서 실험했을때도 activation을 기준으로 weight를 선택하는 것이 중요함을 나타냅니다.

 

이러한 결과를 보고 저자들은 large magnitudes를 갖는 input features가 더 중요하다는 가설을 세웠습니다.

 

Limitation:

* 0.1%의 weight를 FP16으로 유지한 mixed-precision data type은 system implementation을 어렵게 합니다.

* 그렇기 때문에 실제로 FP16을 유지하지 않고 보존할 방법을 찾아야합니다.

 

(c) Protecting Salient Weights by Activation-aware scaling

저자들은 (b)의 mixed-precision으로 인한 하드웨어 최적화 문제를 해결하기 위해  per-channel 스케일링 방법을 제안합니다.

 

Analyzing the Quantization error.

시작하기에 앞서서, weight-only quantization error를 수학적으로 분석합니다.

 

$$y = Wx, \quad y = Q(W)x$$

 

quantization function

$$Q(W) = \Delta \cdot \mathrm{Round}\left( \frac{W}{\Delta} \right), \quad \Delta = \frac{\max(|W|)}{2^{N-1}} $$

 

$N$은 number-of quantization bit $\Delta$는 quantization scaler로, $\max(|W|)$에 의해서 결정됩니다.

 

이제 weight의 한 원소를 $w$라고 할때, 이를 $s > 1$로 scaling하고, x를 $1/s$로 스케일링하면

$$Q(w \cdot s) \cdot \frac{x}{s} = \Delta^{'} \cdot \text{Round} \frac{ws}{\Delta^{'}} \cdot x \cdot \frac{1}{s}$$

 

$\Delta^{'}$는 $s$를 적용한 다음의 새로운 quantization scaler입니다.

 

저자들은 실험을 통해 아래와 같은 요소들을 발견했습니다.

• Round 함수의 평균 오차는 스케일링과 상관없이 [0, 0.5]사이에서 랜덤하게 나왔다.
• single element w에 scaling하는 것은 W의 maximum value에 영향을 주지 않는다. (그러므로, $\Delta' \approx \Delta$로 사용할 수 있다.)
• $\Delta$와 $x$가 FP16으로 되기 때문에, quantization error가 없다.

따라서, quantization error는 아래와 같이 표현할 수 있습니다.

$$\text{Err}(Q(w)x) = \Delta \cdot \text{RoundErr} (\frac{w}{\Delta}) \cdot x$$

 

$$\text{Err}(Q(w \cdot s) \frac{x}{s} ) = \Delta' \cdot \text{RoundErr} (\frac{ws}{\Delta'}) \cdot \frac{x}{s}$$

 

$\Delta' \approx \Delta$ 로 사용할 수 있었기 때문에, $s$가 커질수록 $\text{Err}(Q(w \cdot s) \frac{x}{s} )$는 작아집니다.

 

 

OPT-6.7B 모델에 대해서 다양한 s를 실험한 Table

 

위 Table 2결과를 보면 s=2를 사용했을때의 성능이 제일 좋았습니다.

무조건 $s$를 키우는게 안좋은 이유는, $s$가 커지면 중요하지 않은 weight들의 quantization error가 커지기 때문입니다.

 

 

Searching to scale

모든 weight (중요함/안중요함) 를 모두 고려하기 위해서, input channel별 최적의 scaling factor를 찾는 방법을 사용했습니다.

 

$$s^* = \arg\min_s \mathcal{L}(s)$$

 

$$\mathcal{L}(s) = \left\| Q\bigl(W \cdot \text{diag}(s)\bigr) \cdot \bigl(\text{diag}(s)^{-1} \cdot X \bigr) \;-\; W X \right\|$$

 

$X$는 특정 task에 overfitting되지 않도록 pre-trained dataset에서 작은 calibration set에서 사용합니다.

 

quantization function $Q$가 미분이 되지 않기 때문에, backpropagation을 하지 않습니다. 

* 물론, approximated gradient를 사용하는 연구들도 있지만, 수렴이 불안정한 문제가 있기 때문에 사용하지 않았다고 합니다.

 

중요한 weight channel은 사실 activation scale에 기반한다고 했으니, 더 심플한 search space를 사용할 수 있습니다.

 

$$s = s_X^{\alpha}, \quad \alpha^* = \arg\min_{\alpha} \mathcal{L}\bigl(s_X^{\alpha}\bigr)$$

$s_X^{\alpha}$는 per-channel activation의 average magnitude이며, $\alpha$는 [0, 1]사이에서 grid search로 찾았습니다.

 

 

TinyChat: Mapping AWQ onto edge platforms 

AWQ는 LLM의 memory size를 줄일 수 있지만, 속도 향상으로 이어가는 것은 쉽지 않습니다.

W8A8 quantization (SmoothQuant)는 저장과 계산이 모두 동일한 precision에서 진행되므로, 계산과정이 kernel's epilogue에 통합될 수 있습니다.

*kernel's epilogue?

 

---

 

<접은글>

🌟 kernel’s epilogue란?

먼저 “kernel”이 뭐냐면:

• GPU나 CPU에서 행렬곱/컨볼루션 같은 연산을 수행하는 **하나의 함수(커널)**를 뜻합니다.
  • 예) GEMM kernel (General Matrix Multiply)

이 커널 안에는 보통 크게 두 단계가 있습니다:

1️⃣ Main Loop (Main Computation)

• 행렬 곱셈이나 누산 같은 “핵심 계산”을 반복 수행하는 부분
• 예: 매 tile마다 multiply-accumulate 반복

2️⃣ Epilogue

• “마지막 후처리 단계”입니다.
• 핵심 계산이 끝난 결과에 스케일 적용, bias 추가, activation function, 타입 변환 등을 수행

 

GPT4o 매우 감사.

 

<접은글>

 

---

 

하지만, W4A16의 경우에는 연산전에 W를 16으로 dequantization해야하기 때문에 구현이 어렵습니다.

이러한 문제를 TinyChat이라는 경량화 시스템을 도입해서 해결합니다.

 

 

 

Why AWQ Helps Accelerate On-Devices LLMs

 

위 Figure 3-(a) 에서 generation이 context보다 훨씬 느린 것을 볼 수 있습니다.

 

Figure 3-(c)를 보면 on-device LLM에서 activation에 비해서 weight loading이 무겁습니다.

이는 다른 말로 하면, weight loading을 줄이면 빨라질 수 있습니다.

 

이러한 관점에서 weight-only quantization이 on-device LLM에 특히 유리하다고 이야기합니다.

 

Deploy AWQ with TinyChat

위에서 4bit weight quantization으로 이론상 4배정도의 속도향상이 가능하다는 것을 보았습니다.

GPU에서는 attention, layer normalization, linear projection같은 핵심 연산만 구현에 집중했고, Huggingface FP16대비 3배 이상의 속도 향상을 달성했습니다.

 

그 외에도 On-the-fly dequantization, efficient 4-bit weight packing, kernel fusion을 통해서 추가적인 속도 향상을 가져왔습니다.

 

Experiments

모든 실험에 대해서 group size 128을 사용했습니다.

 

Llama-2와 LLaMA에서 quantization 전과 후의 perplexity를 비교했습니다.

Llama-2와 LLaMA 모두 7B, 13B, 70B 모두에서 RTN, GPTQ보다 더 낮은 perplexity를 보였습니다.

 

Mixtral 모델에도 적용했을때, FP16과 비슷한 수준의 PPL을 달성하며, AWQ가 다양한 아키텍쳐에서도 잘 동작한다는 것을 보였습니다.

 

그 외에 Instruction-tuned, multi-modal language models, visual reasoning, programming and math tasks 등에서도 RTN보다 좋은 성능을 보였습니다.

 

 

2bit quantization을 했을때, RTN은 종종 실패했지만 GPTQ와 AWQ+GPTQ는 성공적으로 양자화가 되었습니다.

(AWQ 혼자 사용하지 않은걸로 봐서는, 2bit on-device에서는 GPTQ보다 낮은 perplexity를 보였을지도 모르겠네요)

 

 

Summary

• 기존 RTN은 perplexity가 높은 문제가 있다.
• 중요한 weight (channel)을 FP16으로 보존하면 quantization error를 낮출 수 있지만, 하드웨어 최적화가 어렵다.
• 추가적인 학습이나 regression없이 scaling을 통해서 중요한 weight를 보호하는 방법을 제안했다.
• 해당 방법은 W4A16으로 weight-only quantization인데, 이러한 방법이 특히 on-device LLM에서 model memory를 줄이면서 weight loading을 빠르게 하며 속도향상이 가능했다.
• 그 외에도 가속화를 위한 TinyChat이라는 platforms를 설계했다.
• 다양한 모델과 task에서 기존 RTN보다 좋고, GPTQ보다 quantization이 잘 되는 것을 보였다.