[논문 리뷰] Mixture of Quantized Experts (MoQE): Complementary Effect of Low-bit Quantization and Robustness

Young J. Kim, Raffy Fahim외 1명|arXiv (Cornell University)|2023. 10. 03.

Multimodal Machine Learning ApplicationsComputer Science인용 수 3

한 줄 요약

이 논문은 Mixture of Quantized Experts (MoQE)를 제안하며, 이는 Mixture-of-Experts (MoE) 모델의 전문가 가중치에만 초저속(2-bit) 양자화를 적용하는 웨이트 전용 양자화 방법이다. MoE 전문가 레이어가 양자화에 대해 내재된 강건성을 지녀, MoQE는 모델 크기를 79.6% 감소시키고, A100 GPU에서 정밀도 손실이 거의 없는 조건에서 1.24배의 추론 속도 향상을 달성하며, Quantization Aware Training (QAT)를 통해 2-bit로 양자화된 경우 밀도 모델을 능가한다.

ABSTRACT

Large Mixture of Experts (MoE) models could achieve state-of-the-art quality on various language tasks, including machine translation task, thanks to the efficient model scaling capability with expert parallelism. However, it has brought a fundamental issue of larger memory consumption and increased memory bandwidth bottleneck at deployment time. In this paper, we propose Mixture of Quantized Experts (MoQE) which is a simple weight-only quantization method applying ultra low-bit down to 2-bit quantizations only to expert weights for mitigating the increased memory and latency issues of MoE models. We show that low-bit quantization together with the MoE architecture delivers a reliable model performance while reducing the memory size significantly even without any additional training in most cases. In particular, expert layers in MoE models are much more robust to the quantization than conventional feedforward networks (FFN) layers. In our comprehensive analysis, we show that MoE models with 2-bit expert weights can deliver better model performance than the dense model trained on the same dataset. As a result of low-bit quantization, we show the model size can be reduced by 79.6% of the original half precision floating point (fp16) MoE model. Combined with an optimized GPU runtime implementation, it also achieves 1.24X speed-up on A100 GPUs.

연구 동기 및 목표

• MoE 모델 추론에서 높은 메모리 대역폭과 지연 시간 오버헤드로 인해 배포 효율성이 제한되는 문제를 해결한다. 이는 하위선형 FLOP 증가에도 불구하고 발생한다.
• MoE 모델의 전문가 레이어가 표준 피드포워드 네트워크(FFN)나 어텐션 레이어보다 저비트 양자화에 더 강건한지 조사한다.
• 모델 품질을 유지하면서 메모리 프로파일을 극적으로 줄이고 추론 속도를 향상시키는 경량의 웨이트 전용 양자화 방법을 개발한다.
• 2-bit로 양자화된 전문가를 가진 MoE 모델이 기계 번역 작업에서 전체 정밀도 밀도 모델을 능가할 수 있음을 입증한다.
• 런타임 엔지니어링을 통해 추론 성능을 최적화하여 현대 GPU 하드웨어에서 측정 가능한 속도 향상을 달성한다.

제안 방법

• MoE 모델의 전문가 FFN 레이어에만 웨이트 전용 양자화를 적용하며, 어텐션 및 밀도 FFN 레이어는 전체 정밀도 그대로 유지한다.
• 전문가 가중치에 대해 2-bit 양자화를 Quantization Aware Training (QAT)와 함께 사용하여 모델 정확도를 유지한다. 반면 3-bit 양자화는 QAT 없이 적용한다.
• 전문가 레이어 가중치가 더 좁고 안정적인 분포를 보이며, 외곽치가 적어 저비트 양자화에 강건함을 관찰하여 이를 활용한다.
• 양자화된 MoE 모델의 추론을 가속화하는 GPU 최적화된 런타임을 구현하여 더 높은 스루풋을 달성한다.
• 다양한 모델 구성 요소(전문가 FFN, 밀도 FFN, 어텐션) 간의 양자화 영향을 비교하여 전문가 레이어의 강건성을 규명한다.
• 다양한 언어 쌍과 모델 크기를 대상으로 WMT 및 내부 검증 데이터셋에서 BLEU 점수를 사용해 양자화 영향을 평가한다.

(a) FFN linear 1 weight distribution across layers

실험 결과

연구 질문

• RQ1 저비트 양자화가 MoE 모델의 다양한 구성 요소, 특히 전문가 FFN 레이어와 밀도 FFN, 어텐션 레이어에 미치는 영향은 어떻게 되는가?
• RQ2 추가 학습이나 校정 데이터 없이 전문가 가중치에 2-bit 양자화를 적용할 수 있으며, 그로 인한 정확도 영향은 어떠한가?
• RQ3 2-bit로 양자화된 전문가를 가진 MoE 모델이 정밀도 손실 없이 전체 정밀도 밀도 모델보다 얼마나 뛰어나게 성능을 냈는가?
• RQ4 전문가 가중치 양자화와 최적화된 GPU 추론의 조합이 모델 스루풋과 메모리 소비에 어떤 영향을 미치는가?
• RQ5 밀도 모델의 표준 FFN 레이어와 비교해 전문가 레이어가 양자화에 대해 얼마나 강건한가?

주요 결과

• MoE 모델의 전문가 FFN 레이어는 밀도 FFN나 어텐션 레이어보다 저비트 양자화에 훨씬 더 강건하며, QAT를 사용할 경우 2-bit 양자화로 인한 성능 저하가 미미하다.
• 2-bit로 양자화된 전문가 가중치를 가진 MoE 모델은 원본 fp16 MoE 모델 대비 모델 크기를 79.6% 감소시킨다.
• 최적화된 GPU 런타임을 통해 MoQE는 원본 fp16 MoE 모델 대비 A100 GPU에서 1.24배의 속도 향상을 달성한다.
• 2-bit로 양자화된 전문가를 가진 MoE 모델은 동일한 데이터셋에서 전체 정밀도 밀도 모델을 BLEU 점수 측정 기준으로 능가한다.
• 밀도 FFN 레이어를 2-bit로 양자화하면 BLEU 점수가 거의 0이 되며, 이는 전문가 레이어가 저비트 양자화에 유일하게 강건함을 확인한다.
• 조금 더 높은 비트 수인 3-bit 양자화라도 WMT 데이터셋에서 평균 BLEU 점수는 1.01점 감소에 그치지만, 2-bit 양자화는 평균 11.54점 감소를 보이지만 여전히 같은 비트 폭에서 밀도 모델보다 훨씬 우수한 성능을 보인다.

(b) FFN linear 2 weight distribution across layers