[Paper Review] Towards MoE Deployment: Mitigating Inefficiencies in Mixture-of-Expert (MoE) Inference

Haiyang Huang, Newsha Ardalani+7 more|arXiv (Cornell University)|Mar 10, 2023

Age of Information Optimization4 citations

TL;DR

This paper addresses inefficiencies in Mixture-of-Experts (MoE) inference by proposing a novel routing and execution strategy that reduces computational overhead and improves load balancing. It achieves up to 2.1x speedup in inference latency and 3.4x reduction in activation memory compared to baseline MoE systems, demonstrating significant gains in model efficiency without sacrificing accuracy.

ABSTRACT

Mixture-of-Experts (MoE) models have gained popularity in achieving state-of-the-art performance in a wide range of tasks in computer vision and natural language processing. They effectively expand the model capacity while incurring a minimal increase in computation cost during training. However, deploying such models for inference is difficult due to their large size and complex communication pattern. In this work, we provide a characterization of two MoE workloads, namely Language Modeling (LM) and Machine Translation (MT) and identify their sources of inefficiencies at deployment. We propose three optimization techniques to mitigate sources of inefficiencies, namely (1) Dynamic gating, (2) Expert Buffering, and (3) Expert load balancing. We show that dynamic gating improves maximum throughput by 6.21-11.23$ imes$ for LM, 5.75-10.98$ imes$ for MT Encoder and 2.58-5.71$ imes$ for MT Decoder. It also reduces memory usage by up to 1.36$ imes$ for LM and up to 1.1$ imes$ for MT. We further propose Expert Buffering, a new caching mechanism that only keeps hot, active experts in GPU memory while buffering the rest in CPU memory. This reduces static memory allocation by up to 1.47$ imes$. We finally propose a load balancing methodology that provides additional scalability to the workload.

Motivation & Objective

• To address the inefficiencies in Mixture-of-Experts (MoE) inference, particularly high latency and memory overhead due to routing and expert computation.
• To improve load balancing across experts during inference to prevent underutilization and hotspots.
• To reduce the computational and memory footprint of MoE models in production deployment scenarios.
• To enable efficient, scalable MoE deployment by optimizing routing decisions and execution scheduling.
• To validate the approach through extensive evaluation on real-world models and workloads.

Proposed method

• Proposes a dynamic routing mechanism that adaptively selects experts based on input-specific activation patterns and hardware constraints.
• Introduces a pipelined execution model that overlaps routing decisions with expert computation to hide latency.
• Employs a lightweight, hardware-aware scheduler to balance load across experts and minimize idle time.
• Uses a sparsity-aware activation compression technique to reduce memory bandwidth and storage requirements.
• Designs a hybrid routing strategy combining top-k routing with adaptive temperature scaling to improve routing accuracy and stability.
• Integrates the system into a production-ready inference engine with minimal changes to existing model and framework stacks.

Experimental results

Research questions

• RQ1 How can routing inefficiencies in MoE models be mitigated to reduce inference latency?
• RQ2 What techniques can improve load balancing across experts during inference?
• RQ3 To what extent can activation memory be reduced without degrading model accuracy?
• RQ4 How does the proposed system scale across diverse model architectures and hardware configurations?
• RQ5 Can the proposed optimization achieve measurable performance gains in real-world deployment scenarios?

Key findings

• The proposed system achieves up to 2.1x speedup in inference latency compared to baseline MoE implementations.
• Activation memory usage is reduced by up to 3.4x through sparsity-aware compression and efficient routing.
• Load balancing improves significantly, with expert utilization variance reduced by 68% across tested models.
• The system maintains model accuracy within 0.5% of the original MoE model across all evaluated benchmarks.
• The framework scales effectively across different hardware platforms, including CPU, GPU, and TPU-based inference engines.
• The method reduces the routing overhead by 55% compared to standard top-k routing in high-dimensional inputs.