EfficientQAT: Efficient Quantization-Aware Training for Large Language Models

Large language models (LLMs) are crucial in modern natural language processing and artificial intelligence. However, they face challenges in managing their significant memory requirements. Although quantization-aware training (QAT) offers a solution by reducing memory consumption through low-bit representations with minimal accuracy loss, it is impractical due to substantial training resources. To address this, we propose Efficient Quantization-Aware Training (EfficientQAT), a more feasible QAT algorithm. EfficientQAT involves two consecutive phases: Block-wise training of all parameters (Block-AP) and end-to-end training of quantization parameters (E2E-QP). To the best of our knowledge, Block-AP is the first method to enable direct training of all parameters in a block-wise manner, reducing accuracy loss in low-bit scenarios by enhancing the solution space during optimization. E2E-QP then trains only the quantization parameters (step sizes) end-to-end, further improving the performance of quantized models by considering interactions among all sub-modules. Extensive experiments demonstrate that EfficientQAT outperforms previous quantization methods across a range of models, including base LLMs, instruction-tuned LLMs, and multimodal LLMs, with scales from 7B to 70B parameters at various quantization bits. For instance, EfficientQAT obtains a 2-bit Llama-2-70B model on a single A100-80GB GPU in 41 hours, with less than 3 points accuracy degradation compared to the full precision (69.48 vs. 72.41). Code is available at https://github.com/OpenGVLab/EfficientQAT.

한 줄 요약

EfficientQAT는 블록별 파라미터 학습과 엔드투엔드 정량화 파라미터 학습을 결합해 대규모 언어 모델의 정량화 성능을 향상시키는 새로운 QAT 알고리즘이다.

핵심 기여도

Block-AP를 통해 모든 파라미터를 블록 단위로 학습하여 정량화 오류를 최소화하고 솔루션 공간을 확장.
E2E-QP를 통해 정량화 파라미터(스텝 사이즈)만 엔드투엔드로 학습하여 하위 모듈 간 상호작용을 고려.
2-bit Llama-2-70B 모델을 단일 A100-80GB GPU에서 41시간만에 학습, 정확도 저하 3점 미만(69.48 vs. 72.41).
7B~70B 규모의 기반, 지시 조정, 멀티모달 LLM에서 기존 정량화 방법 대비 5점 이상 성능 개선.

핵심 아이디어

EfficientQAT는 기존 QAT의 높은 메모리 및 학습 비용 문제를 해결하기 위해 두 단계로 구성된 새로운 알고리즘을 제안한다. 첫 번째 단계인 Block-AP는 트랜스포머 블록 단위로 모든 파라미터(가중치, 스텝 사이즈, 제로 포인트)를 학습하며, 이는 정량화 오류를 줄이고 솔루션 공간을 확장해 정확도를 향상시킨다. 이는 기존 블록 기반 재구성 방법과 달리 추가적인 학습 파라미터를 도입하지 않아도 효과적인 정량화가 가능하다는 점에서 혁신적이다. 두 번째 단계인 E2E-QP는 정량화된 가중치를 고정하고, 스텝 사이즈만 엔드투엔드로 학습함으로써 블록 간 상호작용을 고려해 모델 성능을 추가적으로 개선한다. 이러한 두 단계는 기존 QAT의 정확도와 PTQ의 효율성을 결합한 새로운 접근법이다.

기술적 접근법

**Block-AP**: 트랜스포머 블록 단위로 모든 파라미터(𝐖, s, z)를 학습.
**E2E-QP**: 정량화된 가중치는 고정하고, 스텝 사이즈(s)만 엔드투엔드로 학습.
**하이퍼파라미터**: Block-AP는 4096개 샘플, E2E-QP는 1 에포크로 학습.
**데이터셋**: WikiText2, C4, ReadPajama, MMLU, zero-shot 추론 태스크.
**모델**: Llama-2-70B, 7B~70B 규모의 기반 및 지시 조정 LLM.
**정량화 비트**: 2-bit, 3-bit, 4-bit 등 다양한 수준.

주요 결과

**Llama-2-70B 모델**: 2-bit 정량화 시, 단일 A100-80GB GPU에서 41시간 학습, 정확도 69.48 (기존 72.41 대비 -3점).
**zero-shot 추론 태스크**: Block-AP + E2E-QP 조합 시, 기존 PTQ 방법 대비 약 5점 성능 개선.
**MMLU 정확도**: PEQA 대비 4.5점 상승.
**Perplexity**: Block-AP는 8.53 (기존 BitNet 14.32 대비 38% 개선).

의의 및 한계

EfficientQAT는 대규모 언어 모델의 정량화를 효율적으로 수행할 수 있는 새로운 QAT 알고리즘으로, 기존 QAT의 높은 학습 비용 문제를 해결하면서도 정확도를 유지한다. 특히, Block-AP는 기존 블록 기반 재구성 방법과 달리 복잡한 파라미터 설계 없이도 뛰어난 성능을 보여주며, E2E-QP는 블록 간 상호작용을 고려해 모델 성능을 추가적으로 향상시킨다. 그러나, Block-AP는 학습 샘플 수가 적을 경우 과적합이 발생할 수 있으며, E2E-QP는 스케일이 커질수록 추가 샘플이 필요할 수 있다. 또한, 현재 연구는 주로 2-bit, 3-bit 정량화에 초점을 맞추고 있어, 더 낮은 비트 수에서의 성능은 추가 연구가 필요하다.

실용적 활용

EfficientQAT는 대규모 언어 모델을 저비용으로 정량화하여 모바일, 임베디드 시스템, 클라우드 인프라 등 메모리 제약이 있는 환경에서 활용 가능하다. 특히, Llama-2-70B와 같은 초대형 모델을 2-bit로 압축하면서도 정확도를 유지할 수 있어, 실시간 추론이 필요한 산업 분야(예: 챗봇, 고객 지원 시스템)에 적합하다. 또한, 표준 통일 정량화를 사용하므로 인기 있는 도구박스를 통해 쉽게 배포할 수 있다.