LLM inference은 compute-bound가 아니라 memory-bound이다
즉 파라미터를 메모리에 load하고 store하는 속도가 핵심 bottleneck이다 = memory bandwidth가 latency에 직결됨
"Memory Wall" 이라고 부르기도 함 (https://arxiv.org/pdf/2403.14123)
cf) why memory-bound?
- LLM inference는 대부분 matrix-vector operation이다
- data reuse가 제한적이고, 서로 다른 토큰에 해당하는 벡터를 amortize하기도 힘들다
Quantization?
이를 해결하기 위한 하나의 approach가 quantization이다
어차피 dequantization 및 FP16 computation 연산 비용은 상대적으로 작다, quantizing을 통해 더 작은 memory를 store하고 load하는게 더 이득
performance degration을 최소화하면서 16 또는 32-bit 정밀도의 parameter/weight를 더 낮은 차원으로 quantize
8bit quantize(https://arxiv.org/pdf/2206.01861), GPTQ (4bit), AWQ, SpQR 등등
weight-only quantization이 인기를 끌고 있다 (leave activation full-precision)
1) QAT = Quantization-Aware Training
- 나중에 retraining이 필요, quantization 이후 정확도를 복원하기 위해 weight을 어떻게 adapt할지 추가로 학습해야 함
(low-precision 환경에서 발생할 수 있는 왜곡 및 정보 손실 최소화하도록...?)
- 비록 performance는 더 좋아도, feasible하지 않다. 거대 모델의 학습은 비싸고 해당 학습 환경을 재현하기도 어렵기 때문
2) PTQ = Post-Training Quantization
- retraining이 필요없다, 이미 training 끝난 모델에 quantize 적용
1) Uniform quantization
- weight range를 여러 개의 bin 구간으로 균일하게 분배, 빠른 quantization 연산이 가능하다
2) non-uniform quantization
- 어차피 mem-bound인 llm inference에서는 quantization 연산이 좀 걸려도 상관없다
- 다른 요소를 고려하여 non-uniform하게 weight distribution을 쪼개자 (ex. k-means?)
=> 하지만 여전히 lower bit quantize / smaller model quantize 등등에서 일관적으로 high performance를 얻기는 어렵다
Overview of SqueezeLLM
main bottleneck인 memory issue를 두 가지 전략으로 tackle한다, PTQ 기반이다
weight에만 quantization을 적용하였고, 일반적인 quantize 기법은 layerwise output의 결과를 신경썼지만 여기는 model's final output(=loss)에 최대한 영향이 덜 가는 방향으로 quantize를 했다.
1. sensitivity-based Non-Uniform Quantization
- weight distribution을 non-uniform하게 나눈다
- 단순히 k-means clustering보다 sensitivity 기반으로 가중치를 부여하여 sensitive weight를 더 잘 표현하게 만든다 = weight distribution과 sensitivity를 둘 다 고려 = sensitivity-based weighted k-means clustering
- end-to-end performance가 더 좋아지더라
2. Dense-and-Sparse Quantization
- weight outlier 때문에 quantization 난이도가 올라간다, 소수 때문에 전체가 왜곡될 가능성이 있다
- 따라서 dense 부분과 sparse 부분(outlier)로 나누어 dense part는 더 compact하게 quantize를 하고, sparse 부분은 full precision으로 관리하는 방식을 취한다
=> 기존 다른 접근은 적당히 grouping 하여 간접적으로 해결했지만, 여기서는 directly extract outlier
- perplexity(모델이 데이터를 얼마나 '혼란스러워' 하는지)가 낮아지더라
Approach 1: Sensitivity-Based Non-uniform Quantization
위 사진처럼 weight distribution은 non-uniform한 패턴을 띤다
=> 따라서 약간의 computation overhead가 있더라도 non-uniform quantization을 택하는 게 합리적이다
기존 방식은 주어진 weight distribution를 가장 잘 설명하는 k centroid를 결정하는 k-means clustering을 사용한다
- 3bit quantization이면 k=8
- 실제 weights W와 quantized weights W_Q (k개의 centroid 중 하나) 의 차이를 최소화하는 optimization 문제이다
SqueezeLLM은 sensitivity 기반으로 한 단계 더 발전시켰다
quantization에서의 optimization 문제 = final loss의 perturbation(변동)을 최소화하는 문제
테일러 전개를 할 수 있다, g는 W에서 loss의 gradient, H는 Hessian (2차 미분) 이다
모델이 수렴했다고 가정하면 g는 거의 0에 가깝다고 할 수 있으므로
위와 같은 optimization problem으로 치환할 수 있다.
위 식을 해석해보면 large Hessian values를 가진 weight에 대해서 정밀하게 quantize를 해야 final output의 변동폭이 작아진다
= Hessian value는 weight의 sensitivity이다
= k-means centroid를 정할 때, sensitive한 weight 근처로 value가 위치하도록 해야 한다 !!
연산 및 문제 단순화를 위해 Fisher information matrix F로 근사할 수 있고, 추가적으로 이 F를 diagonal matrix로 근사할 수 있다
최종 식은 weighted k-means clustering 같은 형태가 된다, F_ii가 큰 weight 쪽으로 끌려갈 것.
Approach 2: Dense-and-Sparse Quantization
대부분의 값은 10% 정도에 분포해있으며, 0.01% 정도의 소수의 outlier들이 나머지 90%에 분포해있더라
=> few outlier들을 무시한 상태로 sensitive value들에 집중해보자
W = D + S
- D: dense matrix => sensitivity-based k-means clustering으로 quantize를 진행한다
- S: sparse matrix => outlier를 full-precision으로 보관한다, compressed sparse row(CSR) 포맷으로 관리하자
* 실제 구현에서는
- 0.05%의 layerwise sensitive value를 추출하고
- 0.4%의 outlier를 추출하여
0.45%의 weight를 추출하여 full-precision으로 관리하였다고 함
WX = DX + SX 커널로의 구현?
- DX 계산은 CUDA LUT(LookUp Table) 기반 kernel을 짜서 구현했다
* compressed된 채로 weight을 메모리에서 꺼내서 lookup table에서 대응되는 fp16 벡터를 매핑하여 decompose
* decompose 이후에는 전부 FP16 연산
- SX 연산에서는 balanced hybrid kernel 을 구현했다
* CSR 포맷과 dense activation vector를 연산해야 했는데, CSR이 너무 sparse해서 각 쓰레드가 한 row씩 처리하기에는 너무 놀았다
* 여러 row를 동시에 처리, 대신 synchronize 추가해서
Evaluation
기존 PTQ 기법들과 비교했을 때
- dense-only vs. non-grouped 에서 더 우수했고 (perplexity)
- sparsity level 0.45% vs. grouped 에서도 더 우수한 performance (perplexity)를 보였다
모델의 압축률이 높아질수록 (=low bit precision으로 갈 수록) 기존 PTQ 기법들에 비해 성능 향상 폭이 크더라
instruction following model에서의 quantization을 활용하기 위해 MMLU 벤치마크로 평가해봄
기존 full-precision에 필적하는 정확도를 보이면서 peak memory usage는 훨씬 줄였다
3bit quantized LLaMA 돌리고 CUDA profiling 돌린 결과이다
Full-precision에 비해 약간의 degration(perplexity)은 있지만 latency는 훨씬 낮고 memory 사용량도 훨씬 낮다
=> SqueezeLLM optimally use GPU's memory bandwidth !!
두 가지 기법을 통해 memory wall problem을 효율적으로 해결한 squeeze llm ~
cf) 다른 quantization 논문에서 나온 차트 해석
roofline 모델이 그려지는 이유
- 연산 집약도(x)가 낮은 곳에서는 메모리에서 데이터를 읽어오는 속도가 한계가 됨, 기울기 = 대역폭
- 하드웨어가 낼 수 있는 이론적인 최대 Ops/s에 도달하면 수평선이 됨, compute bound
weight&activation(layer를 통과한 값) quantize를 하면
- 한 번의 연산에 필요한 메모리 접근량이 줄고, 대역폭이 늚. 기울기가 커짐
- peak ops도 늘어남 (fp16 -> int4 연산이 되면서 연산 유닛이 단위 시간 당 머낳은 연산 처리 가능)
weight만 quantize를 하면
- 한 번의 연산에 필요한 메모리 접근량이 줄고, 대역폭이 늚. 기울기가 커짐
- 연산 자체는 dequantize를 해서 fp16으로 이루어진다, peak ops는 그대로이다, 더 빨리 compte bound에 도달?