MLSys 2024Citations: 800+

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

AWQ: LLM 압축과 가속을 위한 활성화 인식 가중치 양자화

Ji Lin, Jiaming Tang, Haotian Tang, et al. (2023)

모든 가중치가 동등하게 중요하지 않으며, 활성화 분포를 기반으로 소수의 핵심 가중치 채널을 식별하여 보호하는 활성화 인식 양자화 AWQ를 제안했다. 이를 통해 GPTQ보다 빠른 양자화와 더 나은 일반화를 달성했다.

배경

LLM의 배포와 서빙을 위해 모델 압축이 필수적이었으며, GPTQ가 사후 학습 양자화의 효과를 입증한 이후 더 효율적이고 하드웨어 친화적인 양자화 방법에 대한 수요가 증가했다. GPTQ는 보정 데이터에 의존하여 가중치를 재조정하므로 과적합 위험이 있고, 양자화 과정 자체가 상대적으로 느렸다. 또한 양자화된 모델의 하드웨어 가속 실행(efficient kernel)에 대한 고려가 부족했다.

핵심 아이디어

AWQ의 핵심 관찰은 가중치의 중요도가 가중치 자체의 크기가 아니라 대응하는 활성화(activation)의 크기에 의해 결정된다는 것이다. 활성화 크기가 큰 채널의 가중치는 전체 출력에 대한 기여가 크므로 양자화 오차에 더 민감하다. AWQ는 이러한 핵심 채널의 가중치를 보호하기 위해, 양자화 전에 채널별 스케일링 팩터를 적용하여 중요한 가중치의 유효 범위를 확대한다. 이 스케일링은 가중치의 양자화 그리드를 조밀하게 만들어 양자화 오차를 줄이며, 역스케일링은 다음 레이어의 활성화에 흡수시킨다. 최적 스케일링 팩터는 보정 데이터에 대한 그리드 탐색으로 결정한다.

방법론

각 레이어에 대해 소량의 보정 데이터(보통 128개)를 순전파하여 채널별 활성화 크기의 평균을 계산한다. 활성화 크기가 큰 상위 1%의 채널을 핵심 채널로 식별한다. 채널별 스케일링 팩터 s를 탐색하되, 스케일링 후 양자화된 출력과 원본 출력 간의 MSE를 최소화하는 s를 그리드 탐색으로 찾는다. 스케일링된 가중치에 대해 표준 그룹 양자화(group-size 128)를 적용한다. 양자화된 모델의 효율적 실행을 위해 W4A16(4비트 가중치, 16비트 활성화) 커널을 구현하여, 메모리 대역폭 병목 상황에서 실제 추론 속도를 향상시킨다.

주요 결과

AWQ는 LLaMA-1/2(7B~70B), OPT(6.7B~66B) 등에서 4비트 양자화 시 GPTQ와 동등하거나 우수한 perplexity를 달성했다. 특히 보정 데이터와 다른 도메인의 평가에서 GPTQ 대비 더 나은 일반화를 보여, 양자화 과정에서의 과적합이 적음을 확인했다. 양자화 속도는 GPTQ 대비 수배 빠르다. 맞춤 W4A16 커널은 FP16 대비 3.2배 메모리 절감과 함께 1.45배의 추론 속도 향상을 달성했다. 지시 미세조정 모델(Vicuna)에서도 양자화 후 다중 도메인 벤치마크 성능이 잘 유지되었다.

임팩트

AWQ는 GPTQ와 함께 LLM 양자화의 양대 표준으로 자리잡았으며, 특히 하드웨어 효율적 배포를 중시하는 산업계에서 널리 채택되고 있다. NVIDIA TensorRT-LLM에 통합되어 상용 추론 서빙에 활용되고 있으며, HuggingFace에서도 AutoAWQ를 통해 폭넓게 지원된다. 활성화 인식이라는 관점은 이후 SmoothQuant, QuIP, AQLM 등 후속 양자화 연구에 중요한 설계 원칙을 제공했다.

관련 Foundation 논문

관련 논문