S wap M o E : Serving Off-the-shelf M o E -based Large Language Models with Tunable Memory Budget (2024)

Abstract: Mixture of experts (MoE) is a popular technique to improve capacity of Large Language Models (LLMs) with conditionally-activated parallel experts. However, serving MoE models on memory-constrained devices is challenging due to the large parameter size. Typical solutions such as memory swapping or expert pruning may lead to significantly higher latency or severe accuracy loss.In this paper, we introduce SwapMoE, a framework for efficient serving of MoE-based large language models with tunable memory budgets. The main idea of SwapMoE is to keep a small dynamic set of important experts, namely Virtual Experts, in the main memory for inference, while seamlessly maintaining how the Virtual Experts map to the actual experts. Experiments have shown that SwapMoE can reduce the memory footprint while maintaining reasonable accuracy. For example, on text summarization tasks with Switch Transformer, SwapMoE can reduce the memory consumption from 14.2 GiB to 4.7 GiB, together with 50% latency reduction and a slight Rouge-2 score drop of 0.041.

## introduction recently, the world has witnessed the great advancement of pre-trained large language models @xcite . such an advancement is driven by the phenomenon that the larger model capacity generally leads to higher intelligence @xcite . among various attempts to scale up neural networks, mixture of experts (moe) @xcite is a promising technique that can avoid linearly increasing computation based on conditional sparse computation. however, memory-costrainted devices is often a major concern in edge ai training and serving (huang et al., 2023b; gim and ko, 2022; wang et al., 2022; ma et al., 2023; kong * work was done while the author was interning at institute for ai industry research (air), tsinghua university. researchers have proposed various methods to reduce the memory footprint of moe model inference. the most straightforward approach is to use memory swapping, i.e. dynamically loading/unloading the parameters from/to the external memory. for example, @xcite propose to load experts into the memory on demand to reduce memory footprint but introduce additional latency overhead, edgemoe @xcite focuses on memory swapping in the decoding phase of moe-based language models and reduces the swapping overhead using quantization. however, it only considers the scenario of single token inference, which does not align with realworld use cases. other algorithm-perspective approaches @xcite propose to prune the experts to reduce model size. however, swapping-based approaches either have to trade latency for a reduced memory footprint, while pruning-based methods would lead to accuracy loss and require model training. how to efficiently serve an off-the-shelf moe model under memory constraints remains challenging. fortunately, there is a unique property in many moe workloads that can be exploited for performance optimization -activation locality. for example, in generative language models, the output tokens are produced one by one, and they mostly belong to the same semantic context. the successive queries of a user are also semantically related. such locality can potentially lead to consistent ac-tivation patterns of moe experts, therefore create room for more intelligent expert management. our work. we introduce swapmoe to enable efficient continuous moe serving under memory constraints. the key idea is to maintain a dynamicallyupdated compact set of virtual experts in the main memory for moe inference, instead of the redundant set of all experts in the original model. the weights of virtual experts are seamlessly updated according to the data distribution locality and profiled hardware capabilities. as such, the memory footprint and latency of each individual inference process in swapmoe are the same as running a smaller moe model, while the advantage of largecapability moe remains since each expert still has the chance to participate in the computation. we implement swapmoe with huggingface transformers @xcite library. to evaluate the effectiveness of swapmoe, we conduct experiments with switch transformer @xcite (switcht for short), and gptsan (gpt, 2023) models on natural language processing tasks, compared with the normal model inference scheme and strong baselines. for example, on text summarization tasks with switch transformer, swap-moe can reduce the memory consumption from 14.2 gib to 4.7 gib (67% less), together with 50% latency reduction and a slight rouge-2 score drop of 0.041, which shows that swapmoe enables the serving of large moe models on resourceconstrained consumer devices with limited memory budgets. we summarize our key technical contributions as follows: ## background and motivation the sparse moe @xcite is the most commonly used, in which only one or few (represented as k) experts are activated for each input. specifically, in a sparse moe with @xmath0, only the i-th expert with max g(x) i is activated for input x, and the output @xmath1 the moe structure is often accompanied by the transformer architecture @xcite , in which the input of each moe layer is a sequence of tokens and each token may choose different experts in one moe layer. ## our design: swapmoe swapmoe utilizes a two-phase holistic design, as shown in figure 4 . in the online phase, the job of swapmoe is to efficiently identify, update, and use a subset of experts (virtual experts) for memoryconstrained moe inference. the job of the offline phase is to obtain an optimal memory plan for the online phase to facilitate efficient and accurate moe inference. ## related work systems optimization for moe model serving. common techniques for optimizing moe model serving include offloading and swapping memory. huang et al. @xcite propose to swap the experts from gpu memory to cpu memory to reduce the memory consumption of moe models, incurring high latency overhead. se-moe @xcite utilizes ring memory offloading to reduce gpu memory consumption of moe models. however, these methods are not suitable for resource-constrained devices with dynamic latency and memory constraints. optimization for dynamic dl model serving. moe is a type of dynamic neural networks. many approaches are introduced to enable or enhance such dynamic dl model serving in general. nimble @xcite is a system that optimizes, compiles, and executes dynamic neural networks on multiple platforms by using a dynamic type system and a lightweight virtual machine runtime. model scaling approaches @xcite @xcite propose to adjust the model size/architecture on edge devices to meet different resource constraints. remix @xcite proposes to use multiple models and dynamically switch between them during the inference process. besides serving dynamic models, researchers have also attempted to slice the static model to dynamic components to achieve different resourceperformance tradeoffs @xcite . as compared with these approaches, our design is fundamentally different because it is based on the unique structure and characteristics of moe. efficient design of moe models. many existing approaches study the efficiency problem of moe models from the model design perspective. gshard @xcite scales up transformer with moe and improves the quality and efficiency of multilingual machine translation. task-moe @xcite extract subnet from a large moe model by a task-level routing strategy. chen et al. @xcite propose to prune non-professional experts according to the downstream task for efficient moe deployment. mpoe @xcite proposes to build a parameterefficient moe architecture by enforcing parameter sharing between the experts. automoe @xcite utilizes neural architecture search to automatically design more efficient moe models. these methods need to modify and retrain the moe model to reduce resource consumption. they are orthogonal to our approach and can use our method to further reduce resource consumption. ## conclusion this paper addresses the challenges of deploying moe models to resource-constrained edge devices. we propose a framework called swapmoe that adaptively reduces the inference costs of moe models based on the memory constraints while preserving the model accuracy. experimental results have demonstrated that our method can significantly reduce the inference costs of moe models with reasonable accuracy degradation. swapmoe creates a nice space of resource-accuracy trade-off of sota large moe models. our work does not have obvious ethical impacts, as we focusing on model inference acceleration.