Toward Efficient Inference for Mixture of Experts

Authors: haiyang  huang
newsha  ardalani
anna  sun
liu  ke
hsien-hsin s lee
shruti  bhosale
carole-jean  wu
benjamin  lee
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 05:40:32

Abstract: Mixture-of-Experts (MoE) models have recently gained steam in achieving the 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 model 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.55$\times$ for LM, 5.75-10.98$\times$ for MT Encoder and 2.58-5.71$\times$ for MT Decoder.
It also reduces memory usage by up to 1.36$\times$ for LM and up to 1.1$\times$ 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 1.47$\times$. Finally, we propose a load balancing methodology that provides additional robustness to the workload. Our code is available at https://github.com/hyhuang00/moe_inference.

Main paper: 
 introduction: 
 A machine learning model's predictive ability increases with the number of parameters. Model capacity has grown at an exponential rate of 10× per year [1] , which in turn has driven demand for computation. Mixture of Experts (MoEs) decouple model capacity from computational demands by using conditionally, sparsely activated neural networks. They can reduce training costs yet improve accuracy [2] [3] [4] for language modeling [5] [6] [7] [8] , machine translation [9] , and image recognition [10, 11] .But training is only half the story. MoE inference is important yet challenging as large language models are deployed for production services. Our experiments show MoE inference is relatively inefficient, requiring much more time to perform the same number of calculations. MoEs are 15× slower for language models and 3× slower for machine translation compared to their FLOP-equivalent dense counterparts. Distillation could shrink models and reduce latency but harm model quality [6, 8, 2] . Optimizations could increase parallelism and GPU utilization, but they narrowly target specific kernels for communication collectives and GPU computation [12, 13] . They lack a comprehensive analysis of inference costs and neglect inefficiencies in the MoE algorithms themselves. We explore optimizations for MoE inference to improve three important dimensions-token throughput, memory use, load balance-without degrading model quality. We begin by identifying sources of inefficiency in MoE inference, breaking down latency and memory use across components of the model architecture. We find the gating function, which assigns tokens to experts, is a major contributor to MoE's high latency and large memory footprint.We address MoE's inefficiency with Dynamic Gating, a new gating function that better matches each expert's computational capacity to its token assignments, thereby reducing communication and computation for placeholders. Dynamic gating reduces latency and memory use while enabling inference with larger batch sizes and fewer GPUs. We implement this new gating function on an open-source MoE Transformer [14] . Dynamic gating could be integrated with other optimizations for distillation, communication collectives, and GPU kernels for even greater benefit. Furthermore, we develop optimizations that allow experts to better use GPU memory and cores. Expert Buffering exploits high temporal locality across experts. It allocates a fixed, but limited, amount of GPU memory for hot, active experts and relies on CPU memory to buffer all other experts. Less frequently accessed experts are brought into GPU memory as needed, significantly reducing demand for GPU memory. Expert buffering is orthogonal to existing memory management techniques such as offloading. Load Balancing mitigates severe load imbalance across experts. Although MoEs are trained with a loss function encouraging load balance, the token distribution during training often differs from that during inference. Our load balancing technique tracks and estimates expert loads (and their hosting GPUs) based on runtime expert activation data. It then redistributes tokens to balance the load, improving system robustness and reducing risks of out-of-memory errors and oversubscribed GPUs.We implement and evaluate these optimizations for language modeling (LM) and machine translation (MT) tasks. These optimizations significantly improve inference throughput, memory use, and load balance. Moreover, they outperform the state-of-the-art and previously proposed optimizations. 
background: 
 In this section, we aims to provide a basic background for readers unfamiliar with MoE transformers. A more detailed description on the MoEs and related works can be found in App. A.Mixture-of-Experts (MoEs) use different models for different inputs to improve versatility and robustness [15] . An MoE consists of multiple, independent models (i.e., experts) and a gating function that assigns inputs to experts. Each input activates only its assigned expert, allowing the model capacity to expand "outrageously" with more tractable increases in computational cost.MoE Transformer. The Transformer model architecture has defined the state-of-the-art for computer vision and natural language processing [16, 17] . Illustrated in Figure 1 , sparse MoE layers replace the FFN block in the Transformer architecture with an MoE block that consists of multiple expert FFNs. These layers use a gating function to decide which experts are most suitable for each token, and then routes tokens to their corresponding experts. Typically, a token is routed to one or two experts in a top-1 or top-2 gating policy. Compared to traditional Transformers, where the FLOP count per batch scales linearly with the number of parameters, MoE networks require much less computation and allow large models to be trained efficiently. MoE Transformers have reduced training costs for large models [2, [6] [7] [8] and achieved high accuracy in vision, text, speech and multi-task learning [18, 5, [19] [20] [21] .Expert Parallelism. MoE models present an interesting trade-off, requiring less computation but more memory usage than traditional Transformers of the same capacity. Expert layers deploy many additional FFNs, which increase model size and demands for GPU memory. GShard [2] addresses these challenges with expert parallelism, distributing workload across multiple GPUs. Each GPU holds a subset of expert FFNs and copy of all other parameters. All-to-all communication is required when distributing tokens to experts and collecting results from experts. 
mixture-of-experts characterization: 
 We characterize MoE Transformers for Language Modeling (LM) and Machine Translation (MT) against FLOP-equivalent dense models. We study recent models on high-performance testbeds, and a detailed description can be found in Table 1 and 2 in App. B. The mini batch size is set to 8 for Language Modeling and 48 for Machine Translation, the largest feasible values under baseline.Under the baseline MoE design, an expert always processes a capacity of tokens regardless of the number of tokens actually assigned. Experts configured with excess capacity will suffer longer latencies and use more memory. But how much of this waste is incurred in practice? We answer this question by analyzing expert activations on several tasks. Language Modeling (LM). We use three domains-Wikipedia, PubMed, Github, from the PILE dataset [22] as input following [8] . Fig. 5 (a) in App. C.1 shows a highly imbalanced load across experts. Multiple hot experts consistently receive a large share of tokens, while other experts consistently receive a small share or no tokens at all. While all inputs show sparse expert activations, the set of hot experts and their hotness levels vary across domains.We evaluate expert activation by performing translation from English to French, Japanese, and Austrian using validation data from NLLB-200 [9] . Fig. 5 (b) in App. C.1 shows that MT also exhibits load imbalance and a small fraction of experts are hotter than others. Noticeably, decoder activation is extremely sparse, and differs from the encoder in one key aspect: it exhibit strong temporal locality. An expert may be active for several consecutive batch, then go inactive again. This temporal locality for hot experts is key motivation for our expert buffering optimization in Section 5.Fig. 7 in App. C.2 examines inference latency. The MoE is 15× slower for language models (LM) and 3× slower for machine translation (MT). While MoE model will be more accurate than the FLOP-equivalent dense model, the MoE will exhibit much higher latencies than the dense. The finding that MoEs perform the same number of FLOPs but require much more time illustrates the need to mitigate the intrinsic inefficiency of MoE inference. Fig. 9 in App. C.2 shows contributors to inference latency for different models and node counts. Prior studies attribute MoE's longer latency to frequent all-to-all communication [2] , which is true for MoE training where computation is distributed across many nodes and all-to-all communication is the main bottleneck. However, our characterization of MoE inference reveals that, while all-to-all communication occurs in multi-node deployments, it is less significant to latency than the computation for the gating function and experts. This finding informs our optimization in Sec. 4.Fig. 8 in App. C reveals contributors to MoE's significant demands for memory capacity. For LM, the dense model only requires 2.2GB on each GPU whereas the MoE requires up to 18.9GB, an increase of 8.6×. For MT, the dense and MoE models use 7.0GB and 21.2GB, respectively, an increase of 3.0×. Beyond the static model parameter usage, MoE models also use much more dynamic memory than dense models. Fig. 10 in App. C illustrates dynamic memory use for the baseline MoE on the Pear cluster. A significant amount of memory is allocated during the gating and reordering computation, and then freed nearly instantaneously. Our close examination of the memory trace indicates that primary cause of this memory use is batch matrix multiplication within the static gating function. We address this challenge with dynamic gating in Section 4.The root cause of MoE's high latency and memory use lies in its static gating policy. Recent implementations assume experts' computational loads are balanced [23, 12, 5, 8] . Under this assumption, token distribution can be simplified and implemented with all-to-all collectives, which we detail later in Fig. 2a (1). Inefficiencies arise when assumptions about load balance fail. If the gating function assigns fewer tokens than an expert's capacity, the remaining capacity is filled with placeholders (i.e., zeros). If more tokens are assigned than the expert's capacity, excess tokens are dropped, with their information retained only by residual connections. Dropping tokens harms accuracy, so capacity C usually set to large values to prevent information loss and accuracy fluctuations. However, large capacities increase latency and memory use. We find that it can leads to a waste as large as 12.8× for LM and 64× for MT (see App. C.3 for detail).We study whether waste during inference is avoidable. If the workload is well balanced and token allocations across experts are comparable, we could reduce waste by simply scaling down expert capacity. On the other hand, if expert activation is sparse, scaling down capacity risks dropping tokens and harming model accuracy. 
dynamic gating optimization: 
 Our characterization reveals a gap between assumptions and practice regarding load balance across experts, which gap leads to poor performance and resource inefficiency during inference. Although prior studies also notice load imbalance across experts [13, 12] , they retain a gating policy that increases expert capacity in response to severe imbalance. Such a policy seeks model accuracy by ensuring the most heavily loaded experts do not drop tokens. But such a policy also exacerbates the high latency and inefficient memory use we have observed.We propose dynamic gating that tunes efficient, variable capacities for experts. This improves upon static gating, which inefficiently sets large, fixed capacities. Changing the gating policy to permit dynamism is non-trivial. Major MoE implementations rely on statically set expert capacities to ensure all messages sent with all-to-all collectives are equally sized [13, 8, 10] . The NCCL all-to-all primitive requires recipients to pre-allocate memory for messages between GPUs. With equally sized messages, each GPU knows the required memory size beforehand. However, when message sizes vary due to dynamic gating and differing token assignments, a lightweight message is needed to notify each GPU recipient of its incoming tensor size.Static Baseline Figure 2a (1) illustrates static gating, which constitutes our baseline. Here, S represents the sequence length, C represents the Capacity Factor, E represents the number of total experts, and D represents the dimension of each token. The gating function generates expert assignments and translates them into E dispatch masks, each with dimension (S, S × C). Entries in the mask are generated as follows. If token i is assigned to expert e and the e-th mask still has capacity, the i-th column of the first empty row is marked 1 and other entries are marked 0. This process produces a highly sparse mask, which is a tensor with dimension (E, S, S × C) that contains at most S 1's. Input tokens are multiplied with the mask to reorder inputs into E sets of inputs, each with S×C tokens, indicating the assignment of tokens to experts. Dynamic Capacity for Gating Figure 2a (2) illustrates our new dynamic gating procedure, which transfers a variable number of tokens to experts and devices. The procedure simplifies token distribution by transforming a vector of gating decisions. First, the procedure performs an argsort to generate indices that sort the decision vector by expert ID. Second, it uses the indices to produce a sorted decision vector. Finally, it counts the number of occurrences of each expert in the decision vector, thereby determining the number of tokens assigned to each expert.Because the number of tokens assigned to each expert varies, dispatch requires two rounds of communication. First, experts are notified about the number of incoming tokens (i.e., size) using an all-to-all collective. This notification happens as soon as sizes are known, allowing its latency to be hidden behind other computation. In parallel with this first round of communication, input tokens are re-ordered with optimized indices and then split based on sizes. Second, the gating function transfers the actual tokens with another all-to-all.After all experts process their assigned tokens, tokens are collected with another all-to-all and sent to their original devices. Tokens are restored to their original order. This re-ordering is typically implemented using batch matrix multiplication but, as in the dispatch stage, the multiplication could be replaced with a more efficient indexing operation.Costs and Benefits Dynamic gating complexity is O(SD + S log S) where S is sequence length and D is token dimension. The dispatch requires a sort of O(S log S), a bin-count of O(S), and an indexing operation of O(SD). The complexity of dynamic gating is much smaller than that of static gating, which requires batch matrix multiplication of O(S 2 EDC) to reorder tokens such that those assigned to the same expert are contiguous. Dynamic gating eliminates the dispatch mask and avoids the multiplication. Instead, it uses an indexing operation that directly places tokens in the desired order.Beyond this complexity analysis, dynamic gating incurs only modest additional communication costs. An additional all-to-all notifies experts about the number of incoming tokens, a single integer that is communicated at very low cost, which only 20 µs on average in our experiments.In exchange for these modest costs, dynamic gating offers several significant benefits. It enhances model robustness by ensuring tokens are not dropped. When an expert receives more tokens, it adjusts its capacity accordingly, preventing load from exceeding capacity and thus avoiding token loss.Additionally, dynamic gating improves computational efficiency by eliminating empty placeholders. When fewer tokens are sent, the expert is notified of the small load, avoiding unnecessary memory allocation and communication bandwidth associated with placeholders. 
expert buffering optimization: 
 Our analysis of expert activation patterns indicates that, although some experts are often inactive, all experts are activated at least a few times across time and batches. This observation motivates our buffering mechanism, which judiciously offloads expert parameters to CPU memory, freeing GPU memory to hold frequently activated experts and enable larger batch sizes.Figure 2b illustrates the buffering mechanism. During inference, each GPU hosts a number of experts and receives their corresponding tokens. An active expert is one that receives tokens for the current batch. If an active expert is not already present in GPU memory, a memory copy transfers expert parameters from CPU memory into GPU memory. The transfer of parameters proceeds in parallel with the transfer of tokens, allowing the buffering mechanism to overlap data movement and hide latency.Scarce GPU memory is managed with an eviction policy that accounts for expert activation patterns. First, the policy evicts the experts that are not active in the current batch as these experts are less likely to be used in the near future (i.e., concept of temporal locality). Second, the policy evicts experts with the Last In, First Out (LIFO) policy.The choice of the LIFO policy is rooted in how recent MoEs have been implemented. When a single GPU hosts multiple experts, the MoE executes experts sequentially by increasing order of their IDs. Suppose the GPU hosts E = 4 experts and caches parameters for 2 experts. Further suppose that the batch activates experts 1, 2 and 3. The MoE Transformer first fetches and executes experts 1 and 2. It then evicts expert 2 and fetches expert 3. By evicting expert 2 instead of expert 1, the cache retains the expert with the shortest re-use distance.No prior work exploits unique MoE characteristics to optimize memory use during inference. As a caching strategy tailored for MoEs, expert buffering is orthogonal to prior memory efficiency schemes and can be seamlessly integrated with other techniques, such as offloading [24, 3] , for even greater memory savings. 
load balancing optimization: 
 Our analysis of token distribution indicates severe computational load imbalance across experts. GPUs that host hot experts can become oversubscribed and vulnerable to out-of-memory errors. GPUs that host cold experts can idle while waiting for others to complete their computation. These observations motivate our load balancing technique, which colocates heavily loaded experts with lightly loaded ones using run-time activation data.Let P mn denote expert placement where m∈{1, . . . , E} is expert ID and n∈{1, . . . , D} is device ID. When P mn = 1, the m-th expert is placed on the n-th device. Let A mb denote expert activation where m is expert ID and b∈{1, . . . , B} is batch ID. Each value in A mb is the fraction of tokens from batch b assigned to expert m.The placement problem can be reduced to the multi-way, number partitioning problem [25] , which is NP-hard. Moreover, this optimization should be constrained such that each GPU hosts the same number of experts. This constraint balances memory use across GPUs and simplifies communication processes. The problem can be formulated as follows.min max m,b n P mn A mb - 1 D subject to m P mn = E D ∀nGreedy Balancing. We implement a greedy algorithm to optimize the assignment of experts to GPUs. The algorithm sorts experts by their average historical load Ãm and sequentially assigns experts to GPUs in descending order. Each step in the sequence assigns an expert to the GPU with the smallest expected load m P mn Ãm . A GPU is excluded from consideration once it reaches designated load. Anti-Correlation Balancing. Although the Greedy algorithm is effective when expert activations are independent (LM, MT-Encoder), it is less effective when activations are correlated (MT-Decoder).With correlated experts, the number of activations that are estimated from historical data A mb is a poor indicator for load. We address this challenge with Anti-Correlation Balancing. Let S ab denote the Pearson correlation between experts a and b that is observed in historical data. We revise the estimate of a GPU's expected load to m P mn ( Ãm + 0.5 * S am ). This revision reduces load and tends to improve expert buffering performance for MT-decoder. 
evaluation: 
 There are few direct baselines for efficient MoE serving. We obtain robust baselines by reimplementing and adapting previously proposed optimizations for inference. We preserve output quality and deliberately exclude token dropping, which harms quality [26] .Methods. Fairseq is our baseline MoE implementation [8] . Tutel improves latency and memory use with custom kernels [13] . It replaces the gating function's sparse matrix multiplication with a hash table lookup to distribute tokens, and creates a custom cumulative summation kernel to reorder inputs. FasterMoE organizes experts into fine-grained groups, combining token communication and expert execution within each group [27] [28] [29] . It overlaps the communication for one group with execution for another group. FasterMoE also places hot experts on each GPU, reducing communication. Megablock uses a custom kernel to execute experts and accelerate operations on a block-sparse matrix, which organizes token inputs and outputs [30, 29] . Nevertheless, the custom kernel depends on features in advanced GPUs, and may compromise backward compatibility. For Megablock, we adapt the experts' MLPs to omit the bias term aligning with Megablock's structural constraints.Clusters. by 1.46×, demonstrating better scalability. We argue that performance under large batch sizes is more important, as small batch sizes are uncommon in real-world inference workloads, particularly in online serving scenarios. Inference with small batch sizes limits throughput due to the reduced computational intensity. With the help of online batching [31] , larger batch sizes can be easily achieved and become more frequent than the small batch size case. We also provide a detailed analysis on the underlying reason for the behavior in App. D.1.Multi-Node. Fig. 4 shows performance when the system scales to two and four nodes. Dynamic gating tunes expert capacity to token load, thereby eliminating wasted communication for placeholders. More efficient communication translates into throughput gains when MoE models are deployed across multiple nodes. Dynamic gating improves throughput by up to 11.55×, 10.98×, 5.71× for LM, MT-Encoder, and MT-decoder, respectively, when compared against Fairseq's static gating.Memory Use. Fig. 11 assesses the impact on gating policy on memory use. Dynamic gating reduces memory use by eliminating the mask for token dispatch and by avoiding computation for placeholders. On the Apple cluster, memory use for activations falls by 79.6%, from 6.29GB to 1.28GB, when performing LM inference with a batch size of 8. Similarly, memory use falls by 44.2%, from 1.89GB to 1.05GB, when performing MT inference with a batch size of 8.When activations use memory more efficiently, the MoE system can support larger batch sizes and achieve higher performance. On the Apple cluster, dynamic gating permits LM and MT to use batch sizes of 64 and 96, respectively. These batch sizes are 8× and 2× than those permitted by Fairseq's static baseline.Cache Miss Rate. Each GPU deploys a cache to hold a subset of the experts assigned to it. When an expert is not found in GPU memory, the system incurs additional latencies to transfer the desired expert from CPU memory to GPU memory. For Apple's MT tasks, we vary cache size from 16 experts-which accommodates all experts, fully occupies GPU memory, and offers zero reduction for static memory use-to 1 expert, which offers a 32% reduction in static memory use.For each cache size, we report the global worst-case cache miss rate. The miss rate is global because it counts misses based on accesses across all GPUs and their memories. The miss rate is worst-case because it reports the highest miss rate across all GPUs and batches. This metric conservatively estimates the least efficient scenario in cache utilization, offering valuable insights for predicting throughput reduction.Fig. 12 (a) indicates cache miss rates for our LIFO policy approximates those from Belady's MIN, the theoretically optimal policy. The figure also indicates that cache miss rates improve most when cache capacity is greater than 5 experts per GPU or 40 experts across 8 GPUs. This observation corresponds with our prior finding that, on average, over 90 experts are not assigned any tokens by the MT-decoder.Throughput. Fig. 4 shows the impact of expert buffering for MT decoding. The cache accommodates 10 experts per GPU and 80 experts across 8 GPUs. This cache size is the point at which Fig. 12(a) indicates diminishing benefits for cache misses.For single-node inference, buffering negates some of the throughput gains from dynamic gating as cache misses impact performance. Nonetheless, dynamic gating and expert buffering together are still competitive with our baselines. For multi-node inference, dynamic gating and expert buffering together report throughput gains of 2.21× and 4.30× over baselines for two and four nodes, respectively.Memory Use. Latency rises as memory use falls. Fig. 11 indicates buffering experts on CPU memory reduces GPU static memory use, for expert parameters, by 2.25GB. But Fig. 13 reports that as cache sizes shrink and memory use falls, latency rises as more experts are transferred between CPU and GPU memories. As expert transfers consume the limited bandwidth between the CPU and GPU, data rates peak at 12 GB/s and increase latency. New technologies that enhance CPU-GPU bandwidth (e.g., NVIDIA's Grace Hopper) can mitigate these latency issues when GPU memory capacity is constrained and caching only a subset of experts is necessary.Load balancing (LB) can further improve latency. This optimization is particularly beneficial when applied to multi-node settings or combined with expert buffering because balance improves cache performance. Note LB is relevant only for dynamic gating where each expert receives a different number of tokens.We analyze load with and without our balancing optimizations using activation data from Section 3.1.We separate this data into two halves: the first half of the activation data is used to generate a device assignment for each expert, and the second half to estimate load under generated assignments. We record Max load, which is the maximum share of the tokens received by a GPU across all batches. Max load is a worst-case scenario and assesses risks from out-of-memory errors. We also record Avg-Max load, which is the maximum share of the tokens received by a GPU averaged over all batches. Avg-Max estimates typical load and assesses performance risks from oversubscribed GPUs.Fig. 14 indicates that Greedy balancing successfully equalizes expert load for LM, reducing loads per GPU from upwards of 0.6 to below 0.4. Fig. 4 shows how balanced load translates into performance. First, Greedy balancing increases throughput by up to 1.11× and 1.19× when compared against pure dynamic gating. Second, it permits larger batch sizes of 64 and 128 when LM is deployed on multiple nodes. Greedy balancing is similarly effective for MT-encoder.Anti-Correlation balancing is robust to MT-decoder's correlated expert activations. It successfully reduces Max and Avg-Max load in most cases, but the balanced load produces only modest throughput gains of 1.02× when compared against pure dynamic gating. 
conclusion: 
 We analyze the behavior of standard MoE Transformer workloads, pinpointing their inefficiencies in inference latency and memory usage. We introduce a Dynamic Gating policy that significantly enhances efficiency in terms of latency and memory demands during inference. Building on this, we propose an Expert Buffering mechanism, demonstrating its effectiveness in substantially reducing memory requirements for MoE inference deployment with a nominal increase in latency. Additionally, we implement load balancing, leveraging historical activation data and heuristic methods to bolster deployment robustness.Recent years witness the boom of MoE LLM models in commercial applications, including API services and personal assistants [32] [33] [34] [35] , and as a result, much more research is needed for efficient inference of MoE models. Techniques such as heterogeneous experts and token-dropping [12, 34] could offer trade-offs between performance and quality, which in turn could be potentially supported by heterogeneous allocations of GPU resources. Heterogeneous experts would motivate even intelligent and dynamic gating functions that assign tokens based on task difficulty. More efficient interconnect (e.g., optical) could benefit expert movement as well as all-to-all communication, which currently remain a challenge. 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: This work focus on the inference efficiency of the MoE model, which has no direct societal impact.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: The paper poses no such risk. Models being used have been open-sourced for a long time.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We cite the original paper that produced the code or dataset used in this work. We used MIT licensed fairseq and the PILE, as well as the CC-BY-SA licensed NLLB dataset.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [Yes] Justification: Our code is well documented with in-line comments.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: No Crowdsourcing experiment conducted.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA] Justification: No Crowdsourcing experiment conducted. No human subjects involved.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: While the MoE Transformer substantially reduces the training cost and FLOPS for large models, the outrageous size of MoE Transformers and the complex expert parallelism [2] poses obstacles for its deployment, including the high GPU memory requirement and the excessive communication overhead of expert assignment. Various approaches have been invented to relieve these obstacles. Switch Transformer [6] and ELSLM [8] use knowledge distillation to distill a large MoE Transformer into a dense model. While distillation reduces the number of parameters, only a small portion (about 30%) of the accuracy gain can be retained. The MoS strategy proposed in DeepSpeed-MoE [12] distills the knowledge to a smaller MoE Transformer with less layers and shared experts. SE-MoE [3] uses pruning to reduce the number of experts in the model.WideNet [36] and MPoE [37] reduce the number of parameters by enforcing parameter sharing.Beyond reducing the parameters, other methods directly reduce computation and communication.The BASE Layer and Switch Transformer also reduce the number of experts each token is assigned to reduce the communication volume and computation. V-MoE [10] further reduces the number by dropping out a large portion of tokens. Hash Layer [5] replaces the gating layer with a precomputed hash function, which reduces the computation cost, but doesn't alleviate the communication overhead.As the MoE Transformer is a type of Transformer, techniques and optimized architectures that enhance Transformer inference speed may apply. Relevant examples include Reformer [38] , Longformer [39] , and Terraformer [40] . However, there is scant discussion of their application to MoE Transformers, and interested readers may find a detailed review in [41] and benchmarks in [42] .In addition to direct modification of model architecture and parameters, deployment strategies, such as offloading strategies and customized kernel functions are also being explored to reduce the GPU resource usage and latency for inference. Offloading and swapping strategies such as [43] swaps unused tensors form the GPU memory to the main memory to reduce the resource requirement. However, existing strategies can only be applied on dense models. Applying these strategies efficiently on conditional neural networks such as MoE is non-trivial, since the data flow graph cannot be constructed in advance due to the conditional computation. FastMoE [27, 28] dissects communication primitives and expert executions on group-basis to overlap these kernels, but it has not been tested on large number of experts. Tutel [13] and DeepSpeed-MoE [12] improve MoE model performance on datacenter-scale systems by combining system and architecture methods with tailored kernels for both Transformer and MoE layers, and specialized communication primitives. The approach combines expert parallelism, model parallelism, and tensor parallelism to significantly boost throughput and reduce latency. However, DeepSpeed-MoE is not designed to conserve GPU resources and therefore may be impractical for many academic users. Megablock [30] utilizes block-sparse matrix to organize token inputs and combines consecutive expert MLPs into a single kernel, but it requires advanced GPU architecture and does not support bias term in expert MLPs. SE-MoE [3] utilizes Ring Memory offloading to reduce GPU usage, achieving better throughput than DeepSpeed-MoE in low-resource scenarios. However, this approach does not leverage expert activation pattern from MoE Transformers. MAD-Max is a performance modeling framework that enables better computation and communication overlapping to improve training and inference throughput, but the work only focuses on traditional MoE-based recommendation systems [44] .Recent years witness the boom of MoE LLM models in commercial applications, including API services and personal assistants [32] [33] [34] [35] . Among these models, two stand out for their notable advancements: the Mixtral series [32] and the DeepSeek-MoE series [34] , both of which merit brief discussion.The Mixtral series reduces the number of experts to eight, while replacing all dense transformer decoder layers with MoE layers. but replaced all dense transformer decoder layer with MoE layer. Such a design allows the model to be deployed on standard server workstations, which are typically equipped with eight GPUs. The small scale of the model ensures the model can be deployed without expert parallelism, placing it outside the scope of this paper.On the other hand, the DeepSeek-MoE further extends the shared expert proposed in [12] , and adopts fine-grained expert separation to mitigate the feature collapse of MoE model [45] expert separation. In this setting, our optimizations for expert buffering and load balancing will reduce memory use and latency for multi-device inference. Using shared experts is tangential and our optimizations could be extended to support them. For example, expert buffering would lock any shared experts into the cache and prevent their eviction from GPU memory to CPU memory.In this section, we provide additional details to the experiments we conducted in this paper. We use the Python native time module to record latency and, separately, PyTorch Profiler to collect detailed traces.Table 1 details model parameters. The table specifies the number of experts (E), how often a FFN is replaced by an MoE (M ), and each expert's capacity fraction (C). We use capacity settings recommended by [8, 9] . the table also details the model architecture and relevant parameters for language tasks.C.1 Expert Activation Fig. 5 and Fig. 6 illustrates the expert activation pattern measured over different workloads. In Fig. 5 , each row representing a batch and each column representing an expert. Colors indicate expert load as a percentage of all tokens in the batch. Hot experts consistently receive a large share of tokens is marked by intense color lines, while other experts receive a small share is marked by lighter color lines. In Fig. 5 (b), we can see that certain experts in both the encoder and decoder receive a large share of all tokens, nearly half of those in the full batch. Many experts receive few tokens. Encoder activation is mostly dense as most experts are almost always activated whereas decoder activation is extremely sparse.Encoder activation patterns are similar across tasks. The encoder captures source language properties (English), which are the same across the three translation tasks. Surprisingly, decoder activation patterns are also similar across tasks despite differences in decoder architectures and target languages (French, Japanese, Austrian).Fig. 7 shows the latency comparison between MoE and FLOP-equivalent dense model. For LM, the dense model requires 74.2ms whereas the MoE requires more than 1.09s. For MT, the dense model encodes and decodes in 101ms and 32ms, respectively, but the MoE requires 2.26s and 90ms.Effects of Model Scaling. We optimize MoE inference based on our characterization of expert computation. Expert activation and locality depends primarily on the number of experts, a value that ranges from 128 to 512 in our studies. If the number of experts remains unchanged, the sparse activations that underpin our optimizations will persist. This is the most likely scenario as MoE models grow by increasing the size of each expert (i.e., feed forward network) or the number of transformer layers rather than the number of experts.Alternative scenarios seem less likely. If the number of experts were to shrink, activations may become less sparse, experts may exhibit greater locality, and our optimizations may be less helpful. However, reducing the number of experts would forgo the accuracy and efficiency benefits of scaling models with MoEs [8] . If the number of experts were to increase, activations may become more sparse and experts may exhibit less locality. This trend would make our optimizations more helpful but seems unlikely as more experts would put further pressure on already scarce GPU memory capacity.Consider a sequence of S tokens for Language Modeling. If the MoE specifies E = 512 experts and C = 0.05 capacity fraction, each expert is configured to process ECS = 512×0.05×S = 25.6S tokens. However, if the MoE implements top-2 gating and each token is processed by two experts, computation is required for only 2S tokens. Thus, experts are configured to compute on many more tokens than they actually receive. This waste factor for LM is 25.6S/2S = 12.8×.Similarly, for Machine Translation, each expert is configured to process ECS = 128×1×S = 128S tokens, but the model requires computation for only 2S tokens due to top-2 gating. The waste factor is 128S/2S = 64×. Such waste indicates MoE models typically perform a large amount of excess computation and communication as well as consume a large amount of extra memory.In this section, we present experimental results that are informative yet could not be included in the main text due to space constraints.In this section, we provide a detailed analysis on the latency measured in Sec. 7.1. Specifically, we compare our method with FasterMoE as well as Megablock on the methodology.FasterMoE implements separate communication and computation kernels, which are executed concurrently to overlap and hide their latencies. This feature of FasterMoE is important for training, which requires significant amounts of time in both kernels. But this feature also incurs significant kernel launch overheads, which are hard to justify for inference, which spends less time in communication.The essential difference between dynamic gating and Megablocks is the difference between several dense matrix multiplications and a single sparse matrix multiplication. Dynamic gating assigns a varied number of tokens to experts and then performs several multiplications on dense matrices. This technique avoids computation on zeros and padding. And computation for dense matrices is more efficient than that for sparse, leading to higher FLOP rates and inference throughput.In contrast, Megablocks concatenates experts into a large matrix and then performs a single multiplication on a blocked compressed sparse row (BCSR) matrix. The BCSR format also avoids computation on zeros and padding, but incurs overheads from matrix metadata (column indices and row offsets) and indexing multiple blocked sub-matrices. The inefficiency of sparse matrix computation grows with matrix size, which in turn grows with batch size. Although dynamic gating requires several dense matrix multiplications, the cost of launching these kernels depends only on the number of experts and is constant as batch size increases.Fig. 11 illustrates the dynamic memory usage of various methods on the Pearcluster and the Applecluster. Memory use for activations is indicated by brighter colors in the figure. Note that GPU memory usage is not solely attributable to expert parameters and activations. Additional factors, such as memory fragmentation and cuBLAS workspaces, may also contribute to memory pressure and out-of-memory errors. But because these elements are beyond the direct control of gating policy, our evaluation focuses primarily on the GPU memory allocated to the MoE model and model activations.Fig. 12 illustrates the cache miss rates under different policy. Figure 13 illustrates how the cache size affects the latency.Fig. 14 illustrates the effect of load balancing on the maximum load. Results are normalized by total batch size, which means numbers represent the share of the total number of tokens each GPU will handle in the batch. We record Max load, which is the maximum share of the tokens received by a GPU across all batches. Max load is a worst-case scenario and assesses risks from out-of-memory errors. We also record Avg-Max load, which is the maximum share of the tokens received by a GPU averaged over all batches. Avg-Max estimates typical load and assesses performance risks from oversubscribed GPUs. Results show our policy successfully reduce risks from out-of-memory errors.We explored several other techniques beyond those already detailed in this paper. Although these explorations did not yield significant results, they may offer researchers insight or cautionary conclusions.Parallel Expert Execution. Figure 9 indicates expert execution is a significant contributor to latency. Considering that each expert's MLP processes only a small fraction of tokens, we attempted to parallelize their execution by creating multiple CUDA streams (e.g., 2 or 4) on each GPU and executing a distinct subset of experts with each stream. Separate CUDA streams parallelize the execution of expert kernels, potentially increasing GPU utilization and reducing latency. Although our experiments did indicate some overlap in expert execution, we found only modest performance benefits accompanied by an increase in memory use due to the additional CUDA streams. Although expert MLPs are relatively small, they are substantial enough to require explicit allocation of GPU resources.Expert Kernel Fusion. In an effort to enhance data locality and minimize kernel launch overheads, we attempted to integrate separate expert MLPs into a unified kernel. We consolidated parameters of the first layer from all expert MLPs into a single tensor and consolidated those in the second layer into another tensor. We then implement a batch matrix multiplication kernel that concurrently executes the first and second layers across all experts. Our experiments did not indicate any performance benefits. Justification: We provide details on the configuration of computer resource being used in the appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).Justification: We conduct research conforming to the Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 


Decision: Paper presents a timely and relevant investigation into optimizing the inference efficiency of MoE models, focusing on language modeling and machine translation tasks. The authors identify key bottlenecks and propose some optimization techniques to address these bottlenecks.

This paper presents valuable contributions to the field of efficient MoE inference, which is rapidly evolving and becoming prevalent both in research, developement but also in deployment (facing users). While some aspects require further refinement and clarification (see below), the work's strengths and potential impact warrant acceptance as a poster presentation.

- Limited Evaluation on Small Batches: One reviewer pointed out the performance drop of dynamic gating compared to Megablock with small batch sizes, which is a relevant concern for certain inference scenarios. While the authors addressed this by highlighting the prevalence of large batches in production settings, further discussion or potential mitigations for smaller batches would strengthen the paper.
- Limited Scope of Evaluation: While the focus on language modeling and machine translation is justified, exploring the impact of the proposed optimizations on other tasks like summarization (as suggested by a reviewer) would broaden the paper's applicability.
- Clarity and Presentation: Some reviewers noted the paper could benefit from improved clarity and presentation, especially regarding the placement of figures and tables.

Reviews: review: paper_summary: paper_summary: This paper analyzed the behavior of standard MoE Transformer workloads and pointed out the bottleneck in inference latency and memory usage. Then it introduces a Dynamic Gating policy instead of static-size computation to improve the efficiency of the gating operation. It also proposes Expert Buffering which offloads inactive experts to CPU and Load Balancing which distributes experts and data in a more balanced way to cooperate with the Dynamic Gating policy to further improve the efficiency. Experiments on LM and MT demonstrate good performance of the proposed method.
strength_weakness: strength_weakness: 1. In Section 4, line 155, the notations S, C, E, D lack explanation when they appear. Although I can guess out their meaning from the context surrounded, it really takes me time to understand these.
2. In Figure 3, it shows that when batch size is less than 32, Dynamic Gating performs worse than Megablock. During inference especially online service, I guess many requests are batchsize=1 and would be better to use Megablock?
questions: questions: N/A
limitations: limitations: N/A
paper_summary: paper_summary: This paper dives deep into the Mixture of Experts architecture, trying to identify its weaknesses and inefficiencies while coming up with novel solutions to improve the architecture in terms of token throughput, memory use and load balance. The authors find out that the static gating function that assigns tokens to each expert contributes most to the high latency and high memory footprint of MoE models. To this end, the authors come up with dynamic gating which improves the token throughput and reduces memory usage. Additionally the authors also add other optimizations such as expert buffering for better usgae of GPU cores and load balancing that takes into account the difference in token distribution at train and inference time. The authors demonstrate the efficacy of their method on Language Modeling and Machine Translation tasks.
strength_weakness: strength_weakness: 1. While there are super detailed experiments on throughput and memory usage and I understand that those are the main talking points of the paper, I would have appreciated if there were some details about how the perplexity or BLEU scores were impacted by adding these optimisations. Some detailed insights into performance which we care about would have been great.

2. The authors have mentioned limitations of their work anywhere.
questions: questions: Check weaknesses.
limitations: limitations: The authors have not a limitations section.
paper_summary: paper_summary: This paper addresses the challenges of efficient inference for Mixture-of-Experts (MoE) models. The authors identify key inefficiencies in MoE inference, particularly in language modeling and machine translation tasks. They propose three main optimization techniques: dynamic gating, expert balancing, and expert load balancing. The authors implement and evaluate these optimizations, demonstrating improvements in inference throughput and memory usage compared to existing methods.
strength_weakness: strength_weakness: - The paper lacks evaluation of existing representative MoE configurations, such as DeepSeek-MoE. Compared to traditional top-1 and top-2 gating, DeepSeek-MoE selects a larger number of experts, which could potentially impact inference performance. This omission limits the comprehensiveness of the study's comparative analysis.
- The proposed method demonstrates throughput advantages over MegaBlock only with large batch sizes. For smaller batch sizes, MegaBlock remains superior. This limitation somewhat restricts the applicability of the proposed method across different usage scenarios.
questions: questions: - Is there a comprehensive experimental comparison available? Figure 3 in the main text only compares the throughput of different methods. However, the latency and memory usage of these methods are not reported. A more complete analysis including these metrics would provide a fuller picture of the proposed method's performance relative to existing approaches.
limitations: limitations: - The placement of most experimental figures in the appendix creates difficulties for readers. While this is understandable given the length constraints of conference submissions, it does impact the paper's readability. If additional space becomes available, the authors should consider including some of the key results in the main text. This would improve the flow of the paper and allow readers to more easily grasp the main findings without constantly referring to the appendix.
paper_summary: paper_summary: The authors propose three techniques to speed up inference of mixtures of experts: 1) dynamic gating: during the all-to-all exchange process, the authors enable sending different number of tokens to each expert, which requires sorting and sending an extra message about the number of tokens; 2) expert buffering optimization: a last-in-first-out scheme to offload experts to CPU memory; and 3) load balancing optimization: greedily selecting experts with the lowest expected load.
strength_weakness: strength_weakness: 1. The techniques are not very substantial. The authors propose a collection of techniques that speed up inference, but the techniques are not thematic. As a result, I feel this paper may be much better suited for systems/industry tracks, where it would be significantly more relevant.
2. The authors do not provide performance analysis (everything is about throughput and latency). From my understanding (which may be wrong), dynamic gating and load balancing should affect performance, and I would be curious to know how it affects that.
questions: questions: 1. The authors show analysis for language modeling and machine translation. How about other tasks, e.g., summarization?
limitations: limitations: Limitations are sparsely discussed throughout the paper.

Comments: comment: Thank you for the detailed response. After carefully reading the reviews and responses, I plan to increase my score to 6.
Thank you for your detailed reply. I'll keep my original score.
Thank You for the clarifications. I would maintain my score of 7.