MiLo: Efficient Quantized MoE Inference with Mixture of Low-Rank Compensators
This addresses the challenge of efficiently deploying large-scale MoE models for AI applications, though it appears incremental as it builds on existing quantization methods.
The paper tackles the problem of significant accuracy loss in Mixture-of-Experts (MoE) models under extreme quantization (e.g., below 4 bits) by introducing MiLo, which uses a mixture of low-rank compensators to recover accuracy with minimal memory overhead, achieving measured latency speedups on 3-bit quantized models.
A critical approach for efficiently deploying Mixture-of-Experts (MoE) models with massive parameters is quantization. However, state-of-the-art MoE models suffer from non-negligible accuracy loss with extreme quantization, such as under 4 bits. To address this, we introduce MiLo, a novel method that augments highly quantized MoEs with a mixture of low-rank compensators. These compensators consume only a small amount of additional memory but significantly recover accuracy loss from extreme quantization. MiLo also identifies that MoEmodels exhibit distinctive characteristics across weights due to their hybrid dense-sparse architectures, and employs adaptive rank selection policies along with iterative optimizations to close the accuracy gap. MiLo does not rely on calibration data, allowing it to generalize to different MoE models and datasets without overfitting to a calibration set. To avoid the hardware inefficiencies of extreme quantization, such as 3-bit, MiLo develops Tensor Core-friendly 3-bit kernels, enabling measured latency speedups on 3-bit quantized MoE models. Our evaluation shows that MiLo outperforms existing methods on SoTA MoE models across various tasks.