[论文解读] Mixture of Quantized Experts (MoQE): Complementary Effect of Low-bit Quantization and Robustness

Young J. Kim, Raffy Fahim+1 more|arXiv (Cornell University)|Oct 3, 2023

Multimodal Machine Learning ApplicationsComputer Science被引用 3

一句话总结

本文提出Mixture of Quantized Experts（MoQE），一种仅对专家权重进行超低比特（2-bit）量化且仅量化权重的量化方法，应用于Mixture-of-Experts（MoE）模型。通过利用MoE专家层对量化固有的鲁棒性，MoQE在不造成显著精度损失的情况下，将模型大小减少了79.6%，并在A100 GPU上实现了1.24倍的推理加速，优于经过2-bit量化且使用量化感知训练（QAT）的稠密模型。

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量化并结合量化感知训练（QAT）以维持模型精度，而无QAT时则采用3-bit量化。
• 利用观察结果：专家层权重具有更窄、更稳定的分布，且异常值更少，从而支持鲁棒的低比特量化。
• 实现一种GPU优化的推理运行时，以加速量化MoE模型的推理，提升吞吐量。
• 对比不同模型组件（专家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 MoE模型中专家层对量化量化的鲁棒性与稠密模型中标准FFN层相比如何？

主要发现

• MoE模型中的专家FFN层对低比特量化显著更鲁棒，与稠密FFN或注意力层相比，使用QAT时2-bit量化仅导致轻微性能下降。
• 与原始fp16 MoE模型相比，使用2-bit量化专家权重的MoE模型实现了79.6%的模型大小缩减。
• 通过优化的GPU运行时，MoQE在A100 GPU上相比原始fp16 MoE模型实现了1.24倍的推理加速。
• 在相同数据集上，2-bit量化专家的MoE模型在BLEU分数上优于全精度稠密模型。
• 将稠密FFN层量化至2-bit会导致BLEU分数接近零，证实专家层在低比特量化下具有独特鲁棒性。
• 即使专家权重使用3-bit量化，WMT数据集上的BLEU平均仅下降1.01分，而2-bit量化仅导致11.54分的平均下降——远优于同比特宽度下稠密模型的表现。

(b) FFN linear 2 weight distribution across layers