기술 요약: MosaicQuant

문제 정의

4비트 양자화는 대규모 언어 모델(LLM)의 메모리 점유율을 줄이고 추론을 가속화하는 데 상당한 잠재력을 가지고 있지만, 비트 폭의 제한으로 인해 발생하는 근본적인 한계에 직면해 있습니다. 즉, 밀집된 일반 값(inliers)과 드물게 나타나는 큰 크기의 값(outliers)을 모두 충실하게 표현할 수 없다는 점입니다. 이러한 불일치는 상당한 정확도 저하를 초래합니다.

기존 솔루션들은 일반적으로 혼합 정밀도 양자화(mixed-precision quantization)를 채택하여, 아웃라이어를 고정밀도(예: FP16 또는 INT8)로 유지하고 나머지는 4비트로 양자화합니다. 그러나 이 방식은 저비트 실행의 균일성을 깨뜨립니다. 고정밀도 폴백(fallback) 연산이 필요함에 따라 추가적인 데이터 이동, 포맷 변환, 정밀도 변환 오버헤드가 발생하며, 이는 결과적으로 저비트 추론이 약속하는 실질적인 하드웨어 가속 효과를 저해합니다.

방법론

저자들은 인라이어-아웃라이어 분리(inlier–outlier disaggregation) 원리에 기반한 통합 4비트 양자화 패러다임인 MosaicQuant를 제안합니다. MosaicQuant는 아웃라이어를 고정밀도로 보존하는 대신, 가중치 행렬 전체를 4비트 구성 요소인 '밀집 베이스(dense base)'와 '희소 잔차(sparse residual)'로 분해합니다.

1. 인라이어–아웃라이어 분리

이 방법은 등가 변환(예: SmoothQuant 또는 QuaRot)을 사용하여 활성화 아웃라이어를 가중치 행렬으로 이동시키는 것부터 시작합니다. 이동된 가중치 행렬 $\hat{W}$는 다음과 같이 분해됩니다:
$$\hat{W} \approx Q(\hat{W}) + R$$
여기서:

• 밀집 4비트 베이스 ($Q(\hat{W})$): 일반적인 인라이어를 충실하게 포착합니다.
• 희소 4-bit 잔차 ($R$): 아웃라이어에 의해 발생하는 양자화 오차를 보상합니다.

결정적으로, 두 구성 요소 모두 4비트로 동작하므로 저비트 실행 파이프라인의 균일성을 유지합니다.

2. 블록 단위 잔차 할당

본 논문은 양자화 오차가 채널 전체에 걸쳐 균일하게 분포하는 것이 아니라, 미세한 블록 수준의 집중(block-level concentration) 현상을 보인다는 점을 식별했습니다. 따라서 MosaicQuant는 잔차 예산을 채널 단위가 아닌 블록 단위로 할당합니다.

비용이 많이 드는 전체 출력 평가 없이 가장 중요한 블록을 선택하기 위해, 저자들은 Hessian-guided $\Delta$-score를 도입합니다. 이 지표는 입력 채널의 활성화 민감도(대각 Hessian 추정치)로 잔차 오차에 가중치를 부여함으로써 각 블록의 출력 왜곡 감소량을 근사합니다. 알고리즘은 가장 높은 $\Delta$-점수를 가진 상위-$K$개 블록을 선택하여 희소 4비트 잔차 보상을 수행합니다. 실험 결과, 단 **10%**의 블록 예산만으로도 출력 왜곡의 50% 이상을 복구하기에 충분하다는 것을 보여주었습니다.

3. ZipperEngine: 통합 실행

희소 잔차를 별도의 커널로 실행할 경우, 활성화를 다시 로드하고 부분 출력을 병합해야 하므로 상당한 지연 시간이 발생하기 때문에, 표현 방식만 통합하는 것으로는 충분하지 않습니다. 이를 해결하기 위해 저자들은 다음과 같은 특화된 실행 엔진인 ZipperEngine을 도입합니다:

• 커널 퓨전(Fuses Kernels): 희소 잔차 계산을 밀집 4비트 GEMM 커널에 직접 통합합니다.
• 중첩된 파이프라인(Overlapped Pipeline): 타일링된 중첩 파이프라인을 활용하여 양자화, 밀집 계산, 희소 보상, 그리고 reduction-add 연산이 Tensor Core, CUDA Core, HBM 상에서 동시에 스케줄링되도록 합니다.
• 결과: 이를 통해 별도의 커널 디스패치 및 중복 메모리 트래픽 오버헤드를 제거하여, 완전한 통합 4비트 추론 파이프라인의 하드웨어 효율성을 실현합니다.

주요 기여

1. MosaicQuant 패러다임: 블록 단위의 인라이어-아웃라이어 분리를 수행하여 정확도 손실을 완화하면서도 완전한 저비트 실행 패턴을 유지하는 4비트 양자화 접근법을 제시합니다.
2. Hessian-Guided 선택: 비용이 많이 드는 정확한 출력 왜곡 평가를 대체하여, $\Delta$-Hessian 점수를 사용해 임계 잔차 블록을 식별하는 효율적인 활성화 인식 블록 선택 전략을 제공합니다.
3. ZipperEngine: 중첩된 파이프라인을 통해 희소 잔차 계산을 밀집 GEMM 커널에 융합하는 밀집-희소 4비트 실행 엔진으로서, 추가적인 잔차 블록이 있음에도 불구하고 실질적인 속도 향상을 가능하게 합니다.

실험 결과

저자들은 LLaMA3 (3B, 8B) 및 Qwen3 (4B, 8B, 14B, 32B) 모델에 대해 MosaicQuant를 평가했습니다.

• 정확도: MosaicQuant는 FP16에 근접한 정확도를 유지합니다. W4A4 양자화 조건에서, MosaicQuant는 다른 저비트 베이스라인들(RTN, SmoothQuant, QuaRot, SpinQuant, ResQ, Atom)보다 성능이 월등히 뛰어납니다. 예를 들어, Qwen3-32B에서 MosaicQuant는 Perplexity 8.4를 달성했으며, 이는 ResQ의 9.2 및 Atom의 9.3과 비교됩니다. 또한 RTN이나 SmoothQuant와 같은 다른 방법들이 Qwen3 모델에서 흔히 겪는 성능 붕괴 현상에서도 견고한 성능을 유지합니다.
• 처리량(Throughput): NVIDIA RTX 4090에서 MosaicQuant는 W16A16 베이스라인 대비 최대 1.24배의 속도 향상을 달릭합니다. 배치 크기 4~32 범위에서 TensorRT-LLM 및 다른 양자화 베이스라인들과 비교하여 일관되게 가장 높은 처리량을 제공합니다.
• 오버헤드: 희소 4비트 브랜치는 최소한의 오버헤드만을 발생시키며, 퓨즈드 파이프라인 덕분에 메모리 사용량 증가는 **8.5~9.5%**에 불과하고 지연 시간 오버헤드는 6.2% 미만으로 유지됩니다.

의의

본 논문은 MosaicQuant가 4비트 LLM 양자화에서의 정확도와 하드웨어 효율성 간의 트레이드오프를 해결한다고 주장합니다. 고정밀도 폴백 연산의 필요성을 제거하고 표현과 실행을 하나의 일관된 저비트 파이프라인으로 통합함으로써, MosaicQuant는 균일한 4비트 실행의 규칙성을 희생하지 않고도 효율적인 LLM 서빙을 가능하게 합니다. 이 접근 방식은 희소한 저정밀도 보상이 밀집 저비트 추론의 처리량 이점을 유지하면서도 아웃라이어로 인한 오류를 효과적으로 복구할 수 있음을 입증합니다.