[논문 리뷰] QMoE: Practical Sub-1-Bit Compression of Trillion-Parameter Models

Elias Frantar, Dan Alistarh|arXiv (Cornell University)|2023. 10. 25.

Machine Learning and Algorithms인용 수 6

한 줄 요약

QMoE는 트리리언 파라미터를 가진 믹스처 오브 응용(MoE) 모델인 SwitchTransformer-c2048와 같은 모델을 1bit 이하로 압축하는 확장 가능한 압축 프레임워크를 도입한다 (파라미터당 0.8비트). 이는 모델 크기를 bfloat16 기준 3.2TB에서 160GB 이하로 줄이며, 정확도 손실을 최소화한다. 커스터마이징된 압축 포맷과 최적화된 GPU 커널을 통해 단일 일반 하드웨어 GPU 서버에서 효율적이고 종단 간(end-to-end) 추론을 가능하게 하며, 런타임 오버헤드는 5% 미만이다.

ABSTRACT

Mixture-of-Experts (MoE) architectures offer a general solution to the high inference costs of large language models (LLMs) via sparse routing, bringing faster and more accurate models, at the cost of massive parameter counts. For example, the SwitchTransformer-c2048 model has 1.6 trillion parameters, requiring 3.2TB of accelerator memory to run efficiently, which makes practical deployment challenging and expensive. In this paper, we present a solution to this memory problem, in form of a new compression and execution framework called QMoE. Specifically, QMoE consists of a scalable algorithm which accurately compresses trillion-parameter MoEs to less than 1 bit per parameter, in a custom format co-designed with bespoke GPU decoding kernels to facilitate efficient end-to-end compressed inference, with minor runtime overheads relative to uncompressed execution. Concretely, QMoE can compress the 1.6 trillion parameter SwitchTransformer-c2048 model to less than 160GB (20x compression, 0.8 bits per parameter) at only minor accuracy loss, in less than a day on a single GPU. This enables, for the first time, the execution of a trillion-parameter model on affordable commodity hardware, like a single server with 4x NVIDIA A6000 or 8x NVIDIA 3090 GPUs, at less than 5% runtime overhead relative to ideal uncompressed inference. The source code and compressed models are available at github.com/IST-DASLab/qmoe.

연구 동기 및 목표

• 트리리언 파라미터를 가진 믹스처 오브 응용(MoE) 모델의 높은 메모리 및 계산 비용을 해결하여 실용적 배포와 연구 접근성을 제한하는 문제를 해결한다.
• 기존의 훈련 후 양자화 기술이 겪는 장벽인 파라미터당 1비트 이하의 정확한 압축을 달성하는 데 도전한다.
• 모델 정확도를 유지하면서도 재훈련이 필요 없는 확장 가능한 압축 파이프라인을 설계하여 소비자 수준 하드웨어에서 효율적인 현장 내 추론을 가능하게 한다.
• 빠른 추론을 위한 압축 포맷과 GPU 최적화된 디코딩 커널을 개발하여 압축된 MoE 모델의 효율적이고 낮은 오버헤드 추론을 가능하게 한다.
• 4x A6000 또는 8x 3090 GPU를 장착한 단일 서버에서 SwitchTransformer-c2048과 같은 거대한 MoE 모델을 실행할 수 있음을 입증한다. 이는 더 넓은 연구 및 배포 가능성을 열어준다.

제안 방법

• 저자체 압축 포맷을 설계하여 저비트 평편도와 가속기 효율성을 고려한, 파라미터당 1비트 이하로 MoE 가중치를 양자화하는 확장 가능한 재훈련이 없는 압축 알고리즘을 제안한다.
• 실시간 압축 해제와 추론을 위한 특화된 GPU 커널을 설계하여, 저비트 양자화된 가중치의 천연 평편도를 활용해 런타임 오버헤드를 최소화한다.
• 압축 포맷과 디코딩 커널을 공동 설계하여 기존 HuggingFace 기반 추론 파이프라인과 GPU 메모리 액세스 패턴에 호환되도록 한다.
• 1.6조 파라미터를 가진 SwitchTransformer-c2048 모델에 압축 파이프라인을 적용하여, 단일 GPU에서 하루 이내에 3.2TB에서 <160GB로 20배의 압축을 달성한다.
• 양자화 과정을 최적화하여 모델 정확도를 유지하며, 사전 훈련 검증 손실과 제로샷 평가 손실에서 미미한 증가만을 보인다.
• HuggingFace 생태계에 통합하여 연구자와 실무자들이 넓게 접근하고 쉽게 배포할 수 있도록 한다.

Figure 1: Example of an MoE Transformer block. Each token is routed to a different fully-connected (FC) block.

실험 결과

연구 질문

• RQ1 트리리언 파라미터를 가진 MoE 모델을 파라미터당 1비트 이하로 압축하면서도 정확도 손실이 심각하게 발생하지 않게 할 수 있는가?
• RQ2 재훈련이 없는 훈련 후 양자화 기술을 사용해 MoE 모델의 1비트 이하 압축을 달성하는 것이 가능한가?
• RQ3 커스터마이징된 압축 포맷과 GPU 최적화된 커널이 소비자 수준 하드웨어에서 압축된 MoE 모델의 효율적이고 낮은 오버헤드 추론을 가능하게 하는가?
• RQ4 압축된 추론의 런타임 오버헤드는 고성능 하드웨어에서 이상적인 압축 해제 실행과 비교해 얼마나 되는가?
• RQ5 최종적으로 압축된 모델을 소비자 수준 GPU를 장착한 단일 서버에서 종단 간(end-to-end)으로 배포하고 실행할 수 있는가?

주요 결과

• 1.6조 파라미터를 가진 SwitchTransformer-c2048 모델은 QMoE를 통해 bfloat16 기준 3.2TB에서 160GB 이하로 압축되었으며, 20배의 압축률과 파라미터당 0.8비트를 달성했다.
• 단일 GPU에서 하루 이내에 재훈련이 없는 훈련 후 양자화 파이프라인을 사용하여 압축이 이루어졌으며, 사전 훈련 검증 및 제로샷 벤치마크에서 정확도 손실이 최소화되었다.
• 압축된 모델의 종단 간 추론이 8x NVIDIA RTX 3090 또는 4x NVIDIA A6000 GPU를 장착한 단일 서버에서 성공적으로 실행되어 소비자 수준 하드웨어에서의 실용적 배포 가능성을 입증했다.
• 압축된 추론의 런타임 오버헤드는 이상적인 압축 해제 실행 대비 5% 미만이었으며, 이는 약 20배 더 많은 GPU가 필요할 것이다.
• 커스터마이징된 압축 포맷과 GPU 커널은 실시간 디코딩을 효율적으로 가능하게 하여 느린 오프로딩을 방지하고 고처리량을 유지했다.
• QMoE 프레임워크는 오픈소스이며 HuggingFace에 통합되어 있어, 대규모 MoE 모델의 연구 및 배포에 있어 접근 장벽을 크게 낮췄다.

Figure 2: Illustration of the offloading execution for the sparse part of a Transformer block. An expert $E_{2}$ and its corresponding input tokens $X_{E}$ are fetched to GPU memory to produce $E^{\prime}_{2}$ , which together with the corresponding outputs $Y_{E}$ are written back to CPU again.