Pre-gated MoE: An Algorithm-System Co-Design for Fast and Scalable Mixture-of-Expert Inference

TL;DR

This work tackles the memory and latency challenges of sparse MoE inference in large language models by proposing Pre-gated MoE, an algorithm-system co-design that decouples expert selection for the next MoE block from the current block's execution. A lightweight pre-gate function deterministically selects the next-block experts, enabling overlap between CPU-to-GPU migration and computation, and offloads the bulk of MoE parameters to CPU memory. Empirical results on SwitchTransformer show substantial improvements: latency reductions of up to ~1.9× over CPU-offload baselines, memory savings of ~4.2× compared with GPU-only, and near-GPU-only throughput with robust accuracy across multiple NLP tasks. The approach enables fast, scalable MoE inference on a single GPU, broadening the practicality of large-scale sparse LLMs.

Abstract

Large language models (LLMs) based on transformers have made significant strides in recent years, the success of which is driven by scaling up their model size. Despite their high algorithmic performance, the computational and memory requirements of LLMs present unprecedented challenges. To tackle the high compute requirements of LLMs, the Mixture-of-Experts (MoE) architecture was introduced which is able to scale its model size without proportionally scaling up its computational requirements. Unfortunately, MoE's high memory demands and dynamic activation of sparse experts restrict its applicability to real-world problems. Previous solutions that offload MoE's memory-hungry expert parameters to CPU memory fall short because the latency to migrate activated experts from CPU to GPU incurs high performance overhead. Our proposed Pre-gated MoE system effectively tackles the compute and memory challenges of conventional MoE architectures using our algorithm-system co-design. Pre-gated MoE employs our novel pre-gating function which alleviates the dynamic nature of sparse expert activation, allowing our proposed system to address the large memory footprint of MoEs while also achieving high performance. We demonstrate that Pre-gated MoE is able to improve performance, reduce GPU memory consumption, while also maintaining the same level of model quality. These features allow our Pre-gated MoE system to cost-effectively deploy large-scale LLMs using just a single GPU with high performance.

Figures (16)

• Figure 1: (a) A dense transformer block that consists of the self-attention layer, feed-forward networks (FFN) layer, normalizations, and residual connections. (b) An MoE block that replaces a conventional transformer block's FFN layer to induce sparsity. The example assumes the MoE block has four expert layers. Each expert has the same dimension as the FFN layer of the corresponding dense transformer block.
• Figure 2: Required number of FLOPs per sequence in deploying SwitchTransformer (MoE) and T5 (dense). In this figure, we show both the "Base" and "Large" model versions of SwitchTransformer and its FLOPs-equivalent T5 (Section \ref{['sect:methodology']} details the model configurations studied in this work). The numbers represent how many experts are available within the MoE block (i.e., dense T5 is equivalent to having just a single expert).
• Figure 3: Memory capacity requirement of deploying SwitchTransformer (MoE) and T5 (dense). MoE parameters include both the expert layer and the gate function, while the rest of the layers are marked as Non-MoE parameters. As depicted, MoE expert parameters account for the majority of the model's memory consumption.
• Figure 4: Pre-gated MoE system for deploying MoE-based LLMs. The memory hungry, sparse MoE parameters are all offloaded to the capacity-optimized CPU memory and are only transferred to the GPU memory when necessary for inference. The rest of the dense, non-MoE parameters are stored locally within the GPU memory.
• Figure 5: Two consecutive MoE blocks employing our proposed pre-gate function. For brevity, we only provide a detailed illustration on the MoE blocks (the rest of the non-MoE layers are consolidated into a single block in this figure). The residual paths within a transformer block are not shown. As depicted, a pre-gate function is trained to select which experts to activate for the next MoE block.