Toward Efficient Inference for Mixture of Experts
Authors: Haiyang Huang
Newsha Ardalani
Anna Sun
Liu Ke
Shruti Bhosale
Hsien-Hsin S. Lee
Carole-Jean Wu
Benjamin Lee
Venue: NeurIPS 2024 poster (2024)
Keywords: Mixture-of-Experts
inference
Primary area: infrastructure
TL;DR: This paper addresses deployment challenges of Mixture-of-Experts models in Language Modeling and Machine Translation by introducing dynamic gating, expert buffering, and load balancing to improve throughput, memory usage, and robustness.

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.

Toward Efficient Inference for Mixture of Experts
Haiyang Huang"
Newsha Ardalani
Anna Sun?
Liu Ke
Hsien-Hsin S. Lee
Shruti 2hosale
Carole-Jean Wu
Benjamin Lee
Duke University
FAIR at Meta
Washington University in St. Louis University of Pennsylvania
Intel Corporation
hyhuang@cs duke. {new, annaysun, shru, aarolejeannu>@meta com
ke ne.1@wustl..m edu lee. sean@gmail com leebcc0seas upenn. edu
Abstract
Mixture-of-Expertt (MoE) models have recently gained steam in achieving the state-of-the-art performance in wide range of tasks in computer vision and natural language processing. They effectively expand the model capacity while incurring 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, provide 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.55x for LM, 5.75- 10. 9898 for MT Encoder and 2.58-5.71 for MT Decoder. It also reduces memory usage by up to 1.36x for LM and to 1.1x x 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.47x. Finally, we propose a load balancing methodology that provides additional robustness the workload. Our code is available at attps:://e //::://gitt github. com/ hyyhangg00mmoe_Bn infiinneeie
Introduction
machine learning model's predictive ability increases with the number of parameters. Model capacity has grown at an exponential rate of 10x per year []] 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 for language modeling g machine translation and image recognition [I0 IT].
But training is only half the story. MoF inference 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. In Section 3 we show that MoEs are 15x slower for language models and 3 slower for machine translation compared 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
to their FLOP-equivalee
Work done while interning/workin/vviiinn FAIR Meta.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

(a)
(b)
(c)
(d)
Tokon
Erbscdings
Exbucding
Eir Entbodingg
Eniraddige
Carling
NIHA
MISA
MB14
Exper
Expeni
Norm
Vorm
Norm
horm
W
W
Gating
Corting
Gating
PTN
Output
Nor
Nor
Ner
Enbosdingg
Embedding
Ernoodding
Figure 1: Illustration of MoE Models. (a) MoE module gating function to assign inputs to experts 15]: (b) Dense Transformer Decoder Layer, which consists of Multi-head Attention (MHA) followed by Feed-Forward Network (FFN); (c) MoE Transformer Layer in which an FFN block replaced by a set of expert FFNs, which operate in parallel; (d) MoE Transformer with expert parallelism. Each device hosts subset of experts. Tokens assigned remote experts are dispatched via all-to-all communication.
lack comprehensive analysis of inference costs and neglect inefficiencies in the MoE algorithms themselves.
We explore optimizations for MoE inference to improve three important dimensions -tokn throughput, memory use, load balance-withowd 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 experts, is major contributor to MoE's high latency and large memory footprint.
We address MoE's inefficiency with Dynamic Gating, new gating function that better matches each expert's 's computational capacity to its token assignments, thereby reducing communication and computation for placeholders. Dynamic gating reduces latency and memory 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.
while enabling
Furthermore, we develop optimizations that allow experts to better use GPU memory and cores. Expert Buffering exploits high temporal locality across experts. allocates fixed, but limited, amount of GPU memory for hot, active experts and relies 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 trained with 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 runtime expert activation data. It then redistributes tokens to balance the load, improving system robustness and reducing risks of out-of-memory errors and oversubscrihed 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.
2 Background
In this section, we aims to provide basic background for readers unfamiliar with MoE transformers. A more detailed description the MoEs and related works can be found in App.

---

Mfxture-of Experts (MoEs) use different models for different inputs to improve versatility and robustness [5]. An MoE consists of multiple, independent models (i.e. experts) and gating function that assigns inputs experts. Each input activates only its assigned expert, allowing the model capacity to expand "outrageously" with more tractable increases in computational cost.
The Transformer model architecture has defined the state-of- -f---r--ttt--tttt-- for computer vision and natural language processing [7]. Illustrated in Figure | 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 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 22 8 and achieved high accuracy in vision, text, speech and multi-task learning
MoE Transformer.
18 5
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-Expert Characterization
We characterize MoE. Transformers for Language Modeling (LM) and Machine Translation (MT) against FL.OP-equivaeent dense models. We study recent models on high-performance testbeds, and detailed description can be found in Table [ and 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.
3.1 Expert Activation
Under the baseline MoE design, an expert always processes 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.
PubMed, Github, from the PILE 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.
Language Modeling (LM). We use three
dataset [22] as input following [8. Fig. S () in App. [h.sshows
Machine Translation (MT). We evaluate expert activation by performing translation from English to French, Japanese, and Austrian using validation data from NLLB-200 [9. Fig. 5])) in App. C. 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
3.2 Latency
Fig. 7 in App. examines inference latency. The MoE is 15x slower for language models (LM) and 3x 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 mitigate the intrinsic inefficiency of MoE inference.
Fig. 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
3

---

MoE training where computation is distributed across many nodes and all-to-all communication is the main bottleneck. However, our characterization of MoË 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.
3.3 Memory Usage
Fig. in App. pppeealees contributors to MoE's significant demands for memory capacity. For LM, the dense model only requires 2.2GB each GPU whereas the MoE requires up to 18.9GB, increase of 8.6x. For MT, the dense and MoE models use 7.0GB and 21. 21.2GB, respectively, an increase of 3.0x. Beyond the static model parameter usage, MoE models also use much more dynamic memory than dense models.
Fig. 10 in App. C iillustrates dynamic memory use for the baseline MoE on the Pear cluster. A sigmificatt 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
3.4 Load Imbalance
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] [2] 5 []. Under this assumption, token distribution can be simplified and implemented with all-to-all collectives, which we detail later in Fig. 2 1). Inefficiencies arise when assumptions about load balance fail. the gating function assigns fewer tokens than 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 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.8x for LM and 64x 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.
4 Dynamic Gating Optimization
Our characterization reveals 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] [2), 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 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 MoF implementations rely on statically set expert capacities to ensure all messages sent with all-to-all collectives are equally sized [13 & 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, 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 dispatch masks, each with dimension (S, C). Entries in the mask generated follows. If token is assigned to expert and the e-th mask still has

---

Decision
1112
Dispatch (ES/S"C))
Recrdered Impul (E(S) CJD)
Irput Takers (30)
Exera Expert Expert
Ratch MarMd
Token All-t- al
OrSEDDC
Multhead
Altontion
Calculation
CPU
(2)
Galing Decision
Rearderod Input (9D)
FFN,
Irput
Size ali-o-eli
Expert Activations
Aegeort
FFN;
Optiral Index
indesing
Check
Cache
FFN,
3/443
0190+
FFN Parameter
Token all-o-al
FFN,
Scried Darision 1
Sine Alltb-ai all
Cache
2:3
3b. Param Copy
Sa. Expert Execution
Token limbedding
(a) Comparison for gating implementation.
(b) Illustration of Expert Buffering.
Figure 2: (a) Comparison between the static gating and our implementation of dynamic gating. We assume E=3, S=6, C=0.5 and top-1 gating in this example. (See Sec. Dec.)( Illustration of the Expert Buffering mechanism. We move the expert parameters CPU memory to reduce burden on GPU memory. (See Sec. 5)
capacity, the i-th column of the first empty row is marked and other entries are marked 0. This process produces highly sparse mask, which is tensor with dimension (E, S, C)) that contains at most 5 I's. Input tokens are multiplied with the mask to reorder inputs into E sets of inputs, each with S SxC' tokens, indicating the assignment of tokens to experts.
Dynamic Capacity for Gating Figure 2a(2) illustrates our new dynamic gating procedure, which transfers 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, uses the indices produce sorted decision vector. Finally, it counts the number of occurrences of each expert in the decision vector, thereby determining the number 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. parallel with this first round of communication, input tokens 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 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 more efficient indexing operation.
Costs and Benefits Dynamic gating complexity is O(SD + S Slo S) where is sequence length and D is token dimension. The dispatch requires sort of O(S log S), bin-coun of O(S), and an indexing operation OSD). The complexity of dynamic gating is much smaller than that static gating, which requires batch matrix multiplication of O 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 indexing operation that directly places tokens in the desired order.
Beyond this complexity analysis, dynamic gating incurs only modest additional communication costs. additional all-to-all notifies experts about the number of incoming tokens, a single integer that communicated at very low cost, which only 20 ps on average in experiments.
5

---

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.
5 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 already present in GPU memory, 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 overlap data movement and hide latency.
Scarce GPU memory is managed with eviction policy that accounts for expert activation patterns. First, the policy evicts the experts that are not active 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 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 and 2. then evicts expert 2 and fetches expert 3. By evicting expert 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 caching strategy tailored for MoEs, expert buffering is orthogonal to prior memory efficiency schemes and can be seamlessly integrated with other techniques, such offloading [24] for even greater memory savings.
6 Load Balancing Optimization
Our analysis of token distribution indicates severe computational load imbalance across experts. GPUs that host hot experts 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 Pmn denote expert placement where me(1.... E}} is expert ID and ne (1....... D} is device ID. When Pmn = 1. the m-th expert is placed on the n-th device. Let Ams denote expert activation where is expert ID and be (1, (l....... B} is batch ID. Each value in Amó 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 hardard. 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
Σ PranA A Amb
subject to Σ Pmnx
Vn D
D
Greedy Balancing. We implement greedy algorithm to optimize the assignment of experts to GPUs. The algorithm sorts experts by their average historical load Ån and sequentially assigns experts to

---

le4
LM-Small Throughput under Different Gating Policy Pear Cluster
Throughout (tokents)
6.
Batch Size
Static
Tutel
Megablock
FastMoE
Dynamic Gating
Figure 3: Throughput Comparison (Pear). Dynamic gating outperforms static, Tutel, and FastMoFE consistently. It outperforms Megablock as batchsize scales. Missing bars represent infeasible combinations of policy and batch size.
GPUs in descending order. Each step in the sequence assigns an expert to the GPU with the smallest expected load Σm Pm Am- A GPU is excluded from consideration once it reaches designated load.
Anti-Correlation Balancing, Although the Greedy algorithm effective when expert activations are independent (LM, MT-Encoder), it is less effective when activations correlated (MT-Decoder). With correlated experts, the number of activations that are estimated from historical data Ame is poor indicator for load. We address this challenge with Anti-Correlatio Balancing. Let Sab denote the Pearson correlation between experts a and b that is observed in historical data. We revise the estimate of GPU's expected load to Σm Pm + 0.5 * Sam). 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 re- implementing 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 S Tutel improves latency and memory use with custom kernels [3] replaces the gating function's sparse matrix multiplication with hash table lookup to distribute tokens, and creates a custom cumulative summation kernel to reorder inputs. FasterMooE organizes experts into fine-grained groups, combining token communication and expert execution within each group 29) It overlaps the communication for one group with execution for another group. FasterMoE also places hot experts each GPU, reducing communication. Megablock uses 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 features in advanced GPUs, and may compromise backward compatibility. For Megablock, we adapt the experts' MLPs omit the bias term aligning with Megablock's 's structural constraints.
our experimental clusters. We use Apple characterize MoE workloads (Table [ and study the impact of our proposed optimizations. Due to limited machine availability on the Apple cluster, we perform additional experiments on Pear with NVIDIA's Ampere. Although our Ampere GPUs provide advanced features for Megablock's 's custom kernels, they offer limited memory capacity and restrict our experiments a single node and smaller LM workloads. We report an average of multiple throughput experiments.
Clusters.
7.1 Impact of Dynamic Gating
Single-Node. Fig. indicate dynamic gating significantly increases throughput for varied batch sizes, tasks, and clusters. LM throughput increases by 6.21x and 3.32x when compared against Fairseq and Tutel, respectively. Similarly, MT-encoder's throughput increases by 5.75x and 5.33x while MT-decoder's 's increases by 2.58x and 1.88x. Beyond throughput increases, dynamic gating permits larger batch sizes by replacing the large dispatch mask with assignment indices and sorted decision vectors.
Compared to other methods, our dynamic gating technique outperforms FasterMoB by up to 2.49x, given the same batch size, by avoiding GPU kernel launch overheads. Moreover, our dynamic gating

---

Throughput under Different Gating Policy on Apple Cluster
Node
165
Threuchput
1.-
2-
(tovaree))
1ed
Encoderr
Node
1e4
MT Ercoder Nodes
1e4
MT Encoder -4 Nodes
15
7.5
10
2-
5.0
2.5
0.0
96
128
256
36
64
96
256
Throughaut bokervs))
163
MT Decoder- Node
1c3
MT Decoder 2 Nodes
163
MT Decoder -4 Nodes
10-
40
60
40
20
8
36
48
64
72
96
8
36
48
64
72
96
128
256
36
48
64
72
96
128
256
Batch
Static TUbel
Dynomic
Dynamic+LB
Dynamic+E8
Dynamic+EBB
Figure Throughput Comparison (Apple). Dynamic gating reduces memory use and communication, enabling larger batches and faster processing compared to static gating. Expert buffering (EB) trades latency for reduced memory use while maintaining high throughput. Load balancing (LB) can further improve latency. Note LB is relevant only for dynamic gating where each expert receives a different number tokens. Missing bars indicate infeasible combinations of policy and batch size.
technique scales to larger batch sizes (and thus higher throughputs) better than Megablocks. Although dynamic gating underperforms for small batch sizes (4-16), it outperforms by increasingly large margins as batch size scales (32-80). When batch size is 80, dynamic gating outperforms Megablock by 1.46x, demonstrating better scalability. We argue that performance under large batch sizes is more important, as small batch sizes are uncommon 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 detailed analysis on the underlying reason for the behavior in App. D.I.
Multi-Node. Fig. shows performance when the system scales to 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.55x, 10.98x, 5.71 x for LM, MT-Encoder, and MT-decoder, respectively, when compared against Fairseq's 's static gating.
Memory Use. Fig. I 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 placcholders. the Apple cluster, memory use for activations falls by 79.6%, from 6.29GB to 1.28GB, when performing LM inference with 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 8x and 2x than those permitted by Fairseq static baseline.
Impact of Expert Buffering
Cache Miss Rate. Each GPU deploys cache to hold subset of the experts assigned to it. When an expert not found GPU memory, system incurs additional latencies to transfer the desired expert from CPU memory to GPU memory. For Apple's 's MT tasks, vary cache size from 16

---

experts whichh which accommodates all experts. fully occupies GPU memory, and offers zero reduction for static memory use-to 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 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 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. 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 falls. Fig. [1 indicates buffering experts on CPU memory reduces GPU static memory use, for expert parameters, by 2.25GB. But Fig. [3 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 necessary.
7.3 Impact of Load Balancing
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 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 generate 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 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 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. shows how balanced load translates into performance. First, Greedy balancing increases throughput by up to 1.11 x and 1.19x 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 '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 x when compared against pure dynamic gating.
8 Conclusion
We analyze the behavior of standard MoE Transformer workloads, pinpointing their inefficiencies in inference latency and memory usage. We introduce 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 MoF inference deployment with a nominal increase in latency. Additionally, we implement load balancing, leveraging historical activation data and heuristic methods bolster deployment robustness.
Recent years witness the boom of MoE LLM models in commercial applications, including API services and personal assistants 32 and as a result, much more research is needed for efficient inference of MoE models. Techniques such as heterogeneous experts and loken-dropping [2] 34 could offer trade-offss 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 henefit expert movement as well all-to-all communication, which currently remain challenge.
Acknowledgement
This work is supported by a National Science Foundation grant CCF-2326606 (Expedition in Com- puting) for the authors at Duke University and University of Pennsylvania. Any opinions, findings, conclusions, or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of this sponsor.
References
[1] Sevilla, J., L. Heim, A. Ho, al. Compute trends across three eras of machine learning, arXiv preprint arXiv:2202.05922 2022.
[2] Lepikhin, D., H. Lee, Y. Xu, et al. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006. 16668, 2020.
[3] Shen, L., Z. Wu, W. Gong, et al. Se-moe: scalable and efficient mixture-of-exxpers distributed training and inference system. arXiv preprint arXiv:2205. 10034, 2022.
[4] Yang, A., J. Lin, R. Men, et al. M6-T: Exploring sparse expert models and beyond. arXiv preprint 1rXiv:215..1508 2021.
[5] Roller, S., S. Sukhbaatar, J. Weston, et al. Hash layers for large sparse models. Advances in Neural Information Processing Systems, 34, 2021.
[6] Fedus, W., B. Zoph, N. Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1-39, 2022.
[7] Du, N., Y. Huang, A. M. Dai, et al. GLaM: Efficient scaling of language models with mixture- of-experts. In International Conference on Machine Learning, pages 5547-5569. PMLR. 2022.
[8] Artetxe, M., S. Bhosale, N. Goyal, et al. Efficient large scale language modeling with mixtures of experts. Y. Goldberg, Kozareva, Y. Zhang, eds., Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 1 11699-11733 11332. Association for Computational Linguistics, Abu Dhabi, United Arab Emirates, 2022.
[9] NLLB Team, M. R. Costa-jussà, J. Cross, et al. No language left behind: Scaling human- centered machine translation. arXiv preprint arXhv:2207.04677 2022.
[10] Riquelme, C., Puigcerver, B. Mustafa, al. Scaling vision with sparse mixture of experts. Advances in Neural Information Processing Systems, 34, 2021.
[11] Xue, F., Z. Shi, F. Wei, et al. Go wider instead of deeper. In Proceedings of the AAAI Conference on Artificial Intelligence, vol. 36, pages 1779-8787 2022.
[12] Rajbhandari, S., C. Li, 7. Yao, et al. DeepSpeed-Mon Advancing mixture-off-expete inference and training to power next gentrreeneatiion ai scale. arXiv preprint arXiv:2201.05596 2022.
10

---

[13] Hwang. C., W. Cui, Y. Xiong, et al. Tutel: Adaptive mixture-of-experts at scale. Proceedings of Machine Learning and Systems, 5:269-287 2023.
[14] Ott, M., Edunov, Baevski, et al. fairseq: A fast, extensible toolkit for sequence modeling. In Proceedings of NAACL-HLT 2019: Demonstrations. 2019.
[15] Shazeer, N., A. Mirhoseini, K. Maziarz, et al. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv:1701.06555 2017.
[16] Vaswani, A., N. Shazeer, N. Parmar, et al. Attention all you need. Advances in Neural Information Processing Systems, 30, 2017.
[17] Dosovitskiy, A., L. Beyer, A. Kolesnikov, et al. An image is worth 16x16 words: for image recognition at scale. In International Conference on Learning Representations. 2021.
Transformers
[18] Hazimeh, H., Z. Zhao, A. Chowdhery, et al. Dselect-k: Differentiable selection the mixture of experts with applications multi-tast learning. Advances in Neural Information Processing Systems, 34:29335 0112333-293477, 2021.
[19] You, Z., S. Feng, D. Su, et al. SpeechMoB: Scaling to large acoustic models with dynamic routing mixture of experts. arXiv preprint arXiv:2105.03003 2021.
[20] Kudugunta, S., Y. Huang, A. Bapna, al. Beyond distillation: Task-level mixture-of-experts for efficient inference. In Findings of the Association for Computational Linguistics: EMNLP 2021, pages 3577-3599. 2021.
[21] Gupta, S., Mukherjee, Subudhi, et al. Sparsely activated mixture-of-experts are robust multi-task learners. arXiv preprint arXiv:2204.07688 2022.
[22] Gao, L., S. Biderman, S. Black, et al. The PILE: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.0002. 2020.
[23] Kim, J., A. Awan, A. Muzio, al. Scalable and efficient moe training for multitask multilingual models. arXiv preprint arXiv:2109 10466, 2021.
[24] Ren, J., S. Rajbhandari, R. Y. Aminabadi, et al. ZeRO-Offoad: Democratizing Billion-Scale model training. In 2021 USENIX Annual Technical Conference (USENIX ATC 21), pages 551-564. 2021.
[25] Graham, R. L. Bounds multiprocessing timing anomalies. SIAM journal on Applied Mathematics, 17(2):416-429 1969.
[26] Fedus, J. Dean, B. Zoph. A review of sparse expert models in deep learning. arXiv preprint arXiv:2209.01667, 2022.
[27] J., J. Qiu, A. Zeng, al. Fastmoe: A fast mixture-of-expert training system. arXiv preprint arXiv:2103. /3262, 2021.
[28] He, J., J. Zhai, T. Antunes, et al. Fastermoe: scale dynamic pre-trained models. In Proceedings of the 27th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming, PPoPP '22, page 120-134. Association for Computing Machinery, New York, NY, USA, 2022.
Modeling and optimizing training of large-
[29] Shoeybi, M.. M. Patwary, R. Puri, al. Megatron-LM: Training multi-billion parameter language models using model parallelism. arXiv preprint
2019.
[30] Gale, T., D. Narayanan, C. Young, et al. Megablocks: Efficient sparse training with mixture-of- experts. Proceedings of Machine Learning and Systems, 5, 2023.
[31] Yu, G.-I., J. Jeong, G.-W. Kim, et Orea: A distributed serving system for (Transformer- Based) generative models. In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), pages 521-538. 2022.
[32] Jiang, A. Q., Sablayrolles, Roux, et Mixtral of experts. arXiv preprint vrXiv:2401.04000 2024.

---

[33] Gemini Team, G. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. Tech. rep.. Google, 2024.
[34] Dai, D., C. Deng, C. Zhao, et al. DeepSeekMoE: Towards ultimate expert specialization in mixture-of-experts language models. arXiv preprint arXiv:2401.06000 2024.
[35] Lin, B., Z. Tang, Y. Ye, al. Moe-llava: Mixture of experts for large vision-language models. arXiv preprint arXiv:2401. 15947, 2024.
[36] Xue, F., Z. Shi, F. Wei, et al. Go wider instead of deeper. In Proceedings of the AAAI Conference on Artificial Intelligence, vol. 36, pages 8779-8787. 2022.
[37] Gao, Z.-F., P. Liu, W. X. Zhao, al. Parameter-efficien mixture-of-expere architecture for pre- trained language models. In Proceedings of the 29th International Conference on Computational Linguistics, pages 3263-3273, 2022.
[38] Kitaev, N., L. Kaiser, A. Levskaya. Reformer: The efficient transformer. International Conference Learning Representations. 2020.
[39] Beltagy, I,, M. E. Peters, A. Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05 2020.
[40] Jaszczur, S., A. Chowdhery, A. Mohiuddin, et al. Sparse is enough in scaling transformers. Advances in Neural Information Processing Systems, 34:9895-9907, 2021.
[41] Tay, Y., M. Dehghani. D. Bahri, al. Efficient Transformers: survey, ACM Computing Surveys, 55(6): 2022.
[42]
FasterTransformer,
Accessed: 2023-
03.
[43] Huang, C.-C., G. Jin, J. Li. Swapadvisor: Pushing deep learning beyond the gpu memory limit via smart swapping. In Proceedings of the Twenty-Fifth International Conference Architectural Support for Programming Languages and Operating Systems, pages 1341-1355. 2020.
[44] Hsia, S., A. Golden, Acun, al. Mad-max beyond single-node: Enabling large machine learn- ing model acceleration on distributed systems. In 2024 ACM/IEEE st Anmual International Symposium on Computer Architecture (ISCA), pages 818 833. IEEE, 2024.
[45] Chi, Z., Dong, S. Huang, al. On the representation collapse of sparse mixture of experts. Advances in Neural Information Processing Systems, 35:34600-34613 2022.
12

---

A Related Work
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-Ne distills the knowledge to a smaller MoE Transformer with less layers and shared experts. SE-MoE
[12]
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 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 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 kemel 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-scal 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, PrepSpeed-Moke 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 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
Recent years witness the boom of MoE LLM models in commercial applications, including API services and personal assistants B2
35 Among these models, two stand out for their notable advancements: the Mixtral series 32 and the DeepSeek-MofF 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-Mon MoE further extends the shared expert proposed in 12). and adopts fine-grained expert separation to mitigate the feature collapse of MoE model [45). Our optimizaaions are relevant for DeepSeek-MofF because, even when the MoE layer activates multiple experts (e.g.. six in 16B model), many of MoE's inefficiencies remain. Expert sparsity remains problem because the total number of experts is large, especially due to DeepSeck-MoE's approach to fine-grained
13

---

Task
Type
Size
E
M
C
Dense
125M
LM-Small
MoE
15B
512
2
0.05
Dense
355M
LM
MoE
52B
512
2
0.05
Dense
3.3B
MT
MoE
54.5B
128
4
Task
Type
Layers
TD
HD
Vocab
LM-Small
Dense/MoE
768
3072
51200
LM
Dense/MoE
24
1024
4096
51200
MT
Dense/MoE
48
2048
8192
256206
Table 1: Model detail for Language Modeling-Smin Small (LM-Small),, Language Modeling (LM) and Machine Translation (MT). Hyperparameters include number of experts (E), MoE layer interval (M), and expert capacity fraction (C). Model parameters include token dimension (TD), hidden dimension (HD), and vocabulary size.
Cluster
Apple
Pear
2xtelel Xeon E5-2698 v4
2 Intel Xeon Gold 5317 with 64GB memory
CPU
with 700GB memory
CPU-GPU
16GB/s via PCle 3.0
32GB/s via PCIe 4.0
8xNVIDIA Tesla V100,
4xNVIDIA RTX A5000, with 24GB memory
GPU
with 32GB memory
connected by NVLink
Nodes
Single Node and Multi Node Single Node only
Table 2: Experimental Clusters.
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.
B Experimental Details
In this section, provide additional details the experiments conducted in this paper. We use the Python native time module to record latency and, separately, PyTorch Profiler to collect detailed traces.
Table detaill model parameters. The table specifies the number of experts (F), how often a FFN replaced by an MoE (M), and each expert's capacity fraction (C". We use capacity settings recommended by [ 9]. the table also details the model architecture and relevant parameters for language tasks.
C Workload Characterization Details and Additional Discussion
C.1 Expert Activation
Fig. and Fig. illustraaes the expert activation pattern measured over different workloads. In Fig. 5 each row representing batch and each column representing an expert. Colors indicate expert load as 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 small share marked by lighter color lines.
14

---

Expert Activation Pattern Selected Layers and Tasks
Wikipedia: Layer
PubMed Abstracts: Layer 5
GitHub: Layer
activated tokens
0.01
0.00
192
288
480
0
96
192
288
480
96
1288
384
480
(b)
fra: Encoder Layer 2
jpn: Encoder Layer 2
ast: Encoder Layer 2 % activated tokens
0.10
0.05
120
48
72
120
24
72
96
120
fra: Decoder Layer 4
jon: Decoder Layer 4
ast: Decoder Layer activeted
0.50
0.00
120
0
24
120
24
96
120
Expert ID
Expert ID
Expert ID
Figure 5: Expert activation patters on selected (a) LM and (b) MT layers. Activation reported
percentage of all tokens in batch. Activations exhibit significant, consistent imbalance on all tasks.
Average Number Inactive Experts
LM 512 Experts
MT Encoder 128 Experts
Decoder 128 Experts
80
10
Unactvatad Exparts.
Wikipedia
eng-ast
102
60-
eng-ipn
GitHub
eng-fra
96
94
MoE Layer ID
MoE Layer
McE Layer
Figure 6: Average number inactive experts for (a) LM, (b) MT encoder, (c) MT decoder. LM and MT' encoder activate most, if not all, experts. MT decoder exhibits extremely sparse activations.
In Fig. b]b), we can see that certain experts in both the encoder and decoder receive large share of all tokens, nearly half of those in the full batch. Many experts receive few tokens. Encoder activation mostly dense as most experts are almost always activated whereas decoder activation is extremely sparse.
Encoder activation patterns 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).
C.2 Latency and Memory Breakdown
Fig. 7 shows the latency comparison between MoE and FL.O--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 the number of experts, value that ranges from 128 to 512 our studies. If the number experts remains unchanged, the sparse activations that underpin 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.
optimizations will persist. This is the most likely scenario MoE
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.
15

---

2.26
MoE
2
Latency (s)
Dense
1.09
1
0.07
0.10
0.09
0.03
LM
MT Encoder
MT Decoder
Models and Tasks
Figure 7: Inference latency for MoE and Dense models on single node. MoE is slower by 15x for LM, 22x for MT encoding, and 3 for MT decoding.
Memory Consumption of MoE Model and Dense Counterpart
30
Embedding Table (Static)
Attention (Static)
Other layers (Static)
Memory Consumption (GB)
25
Expert Parameters (Static)
Activation (Dynamic)
20
15-
10
5
MT MoE
MT Dense
LM MoE
LM Dense
Figure 8: Inference memory use for MoE and dense models. MoEs use more memory for expanded model capacity (multiple experts) and model activations. Results shown using batch sizes and 48 for LM and MT, respectively.
Latency Composition of MoE Model
LM 1 Node
LM Nodes
Gating & Reorder
LM Nodes
Expert Execution
MT 1 Node
All2all Communication
MT 2 Nodes
Other layers
MT 4 Nodes
20
80
Latency Percentage
Figure 9: Inference latency for MoE models. Beyond communication, the gating function and expert
execution are significant contributors latency.
16

---

Memory Allocated (GB)
Gating and Reordering
12
10
Expert Execution
8
100
200
300
400
500
Time (ms)
Figure 10: Memory trace of bascline MoE implementation Pear cluster. Gating and reordering briefly allocates large amount of memory.
However, reducing the number of experts would forgo the accuracy and efficiency benefits of scaling models with MoEs [8] the number of experts were to increase, activations may become more sparse and experts may exhibit less locality. This trend would make optimizations more helpful but seems unlikely as more experts would put further pressure on already scarce GPU memory capacity.
C.3 Calculation on Waste Factor
Consider a sequence of tokens for Language Modeling. the MoE specifies = 512 experts and C 0.05 capacity fraction, each expert is configured to process ECS = x0.2xx 25.6S tokens. However, if the MoE implements top-2 gating and each token processed by two experts, computation is required for only 25 tokens. Thus, experts are configured to compute many more tokens than they actually receive. This waste factor for LM 25.6S/2. 2S 12.8x.
Similarly, for Machine Translation, each expert is configured to process ECS = 128 x x 1285 tokens, but the model requires computation for only 25 tokens due to top-2 gating. The waste factor is 128$/2S = 64x. 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.
D Additional Experimental Results and Analysis
In this section, we present experimental results that are informative yet could be included in the main text due to space constraints.
D.1 Analysis Latency
In this section, provide detailed analysis the latency measured in Sec. 7.1 Specifically, compare our method with FasterMoF as well as Megablock on the methodology.
FasterMoE implements separate communication and computation kernels, which are executed concur- rently to overlap and hide their latencies. This feature of FasterMoE 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 single sparse matrix multiplication. Dynamic gating assigns 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 more efficient than that for sparse, leading higher FLOP rates and inference throughput.
17

---

In contrast, Megablocks concatenates experts into a large matrix and then performs single multiplica- tion 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.
D.2 Memory Usage
Fig. II illustrates the dynamic memory usage of various methods on the Pearcluster and the Ap- plecluster. 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.
D.3 Additional Figures and Results for Effect of Expert Buffering
Fig. 12 illustrates the cache miss rates under different policy. Figure illustrates how the cache size affects the latency.
D.4 Additional Figures and Results for Effeet of Load Balancing
Fig. 14 iillustrates the effect of load balancing 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 GPU across all batches. Max load is worst-case scenario and assesses risks from out-of out-of-memory errors. We also record Avg-Max load, which is the maximum share the tokens received by GPU averaged over all batches. Avg-Ma Max estimates typical load and assesses performance risks from oversubscribed GPUs. Results show our policy successfully reduce risks from out-of of-memory errors.
E
Other Potential Optimizations
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 conclu- sions.
Parallel Expert Execution. Figure 9 indicates expert execution is a significant contributor to latency. Considering that each expert's 's MLP processes only small fraction of tokens, we attempted to parallelize their execution by creating multiple CUDA streams (e.g., 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 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 integrate separate expert MLPs into unified kernel. We consolidated parameters of the first layer from all expert MIPs into single tensor and consolidated those in the second layer into another tensor. We then implement batch matrix multiplication kernel that concurrently executes the first and second layers across all experts. Our experiments did not indicate any performance benefits.
18

---

LM Memory Consumption under Different Gating Policy on Apple Cluster
Memory Consumption (GB)
25.0
Static Gating Policy
27.5
Gating Policy
Dynamic Gating Policy
18.0 10.0
17.5
15.0
12
12.0
132
13.2
13.9 13.9
12.5
32
64
Batch Size
MT Memory Consumption under Different Gating Policy on Apple Cluster
26
Static Gating Policy
Memory Consumption (GB)
24
Tutel Gating Policy
Dynamic Gating Policy
Dynamic Gating Policy Expert Buffering
222
22.4
20
18
16-
16
14
36
48
64
72
96
Micro Batch Size
LM-Small Memory Consumption under Different Gating Policy Pear Cluser
20
Static Gating Policy
Memory Consunption (GB)
18
Tutel Gating Pollicy
Megablock
11
16
Dynamic Gating Policy
12
8-
Micra Batch Size
Figure 11: Memory Use Comparison (Apple, Pear). Dynamic gating reduces memory use, enabling larger batches. Expert buffering further reduces memory use for model parameters. Light shade is dynamic memory for activations. Dark shade is static memory for model parameters. Missing bars indicate infeasible cases.
19

---

Worst-Case Cache Miss Rate
1.0
Cache Miss Rate
LIFO
0.5
LIFO Anti-corr
Belady's MIN
0.0
12
Cache Size per Device
Worst-Case Cache Miss Rate Compared Against Belady's MIN
Cache Miss Rate Diff
LIFO
0.05
LIFO w/ Greedy
LIFO w/ Anti-corr
Belady's MIN
0.00
2
4
6
10
12
14
16
Cache Size per Device
Figure 12: Cache Miss Rates. For trace of MT-decoder's expert activations, (a) misses with and without load balaneing and (b) misses compared against those from Belady's MIN. Miss rates are further reduced by load balancing.
Cache Latency v.s. Memory Usage
Decoder Latency (s)
0.4
4
0.3
6
8
0.2
10
12
0.1
14
16
0.0
17
18
19
20
21
22
Memory Usage (GB)
Figure 13: Cache Sizes vs Latency. For MT-decoder, decreasing the cache size decreases memory usage but increases latency. Cache size (reported number of experts) per GPU marked on plot.
Load per device w/ or w/o optimizations
Load device
Wiki
PubMed Abstracts
GitHub
1.0
Max Load: Naive
Max Load: Greatly
0.5
Avg car:: Gowedy
0.0
10
10
10
MoE Layer ID
MaE Layer ID
MoE Layer ID
eng-fra Encoder
eng-jpn Encoder
eng-ast Encoder
0.50
Gready
0.25
Ave Max Load: Naive
Avg Nox Load: Greedy
Load
0.00
McE Layer
MoE Layer
Layer
device
eng-fra Decoder
eng-jpn Decoder
eng-an ast Decoder
Max Load: Gready
Max Load: Anti-comr
Aug Load: Groady
Avg Nuo Loud: Anti-cor
MoE Layer ID
MoE Layer ID
MoE Layer ID
Figure 14: Load Balancing. Greedy and Anti-correlation algorithms balance load and reduce maximum load across GPUs, reducing risks from out-of out-of-memory memory errors poor performance from oversubscribed devices.
20

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We provide solid experimental results on multiple computing clusters to support the contribution stated and the scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss the limitations of our method in certain less likely scenarios, including when applied to larger model scales, small batch sizes, and varied tasks, as detailed in the Methods and Evaluation sections
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
21

---

Answer:
[NA]
Justification:
This paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: We disclose the publicly available dataset, model checkpoints and settings needed to reproduce our results.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
22

---

reproducible, they Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: The dataset and model checkpoint used in our paper is open-sourced and can be found online. Our code will also be open-sourced and is currently waiting for approval.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://niis.co cc/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/
Coaee//ddeeeebmmmcia ssi onPolicy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Details of our experiments provided in the characterization well the evaluation section.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification:
We performed multiple rounds of experiments and tested the result.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
23

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma enror bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
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 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).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics atttps://neurrs http.///pcbccctctt.ccccccee.....
Answer: [Yes]
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 cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. 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 why the paper does not address societal impact.

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess 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 direct path to any negative applications, the authors should point it out. For example, it is legitimate to point that 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 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 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 attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
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 (c.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper poses such risk. Models being used have been open open-sourced for long time.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have 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 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 best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We cite the original paper that produced the code dataset used in this work. We used MIT licensed fairseq and the PILE, as well the CC-BY-SA licensed NLLB dataset.
Guidelines:
The answer NA means that the paper does 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 used and, 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.
25

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssi
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 not available online, the authors are encouraged to reach out to the asset's s creators.
13. 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/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent 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.
14. Crowdsoureing 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 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 the supplemental material is fine, but if the main contribu- tion 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.
15. 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 equivalent approval/review based the requirements of your country 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 research with human subjects.
Depending 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.
26

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere 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.
27


Reviews:
[{"review_id": "p7ADDxdU8g", "reviewer": "Reviewer_BGdS", "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.", "questions": "N/A", "limitations": "N/A", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The authors really do very detailed systematic and sound analysis of the bottlenecks in MoE and point out the modules that affect the efficiency, which is a strong motivation of the proposed method.\n2. In view of the inefficiency of the gating operation, the authors proposed a dynamic gating policy, along with expert buffering and load balance to systematically improve the efficiency.\n3. The experiments are convincing to address the motivation and support the methods.", "weaknesses": "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.\n2. 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your enthusiastic and encouraging review of our work. Below are our responses towards each point and question raised in the review:\n- **The meaning of notation S, C, E and D in Section 4.** \n\nWe appreciate the reviewer pointing out that these notations are not clearly defined when they appear. 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. In future revision, we will add definitions to these notations in both Section 4 and Figure 2 to enhance clarity.\n- **Dynamic Gating performs worse than Megablock when the batch size is less than 32.** \n\nPlease check our common response 1 for more details.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2hJIzk4mAy", "reviewer": "Reviewer_dmgJ", "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.", "questions": "Check weaknesses.", "limitations": "The authors have not a limitations section.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. I believe the paper is expertly written and provides excellent motivation for coming up with methods for making inference of MoE models efficient. I especially enjoyed reading Section 4 wherein each and every design choice is explained in detail. \n\n2. The optimisations introduced in the paper are quite novel and lead to insane gains in throughput and memory usage. These optimisations are simple and can be easily used in practice for serving large MoE models. The results are pretty strong as well in terms of speedups gained.", "weaknesses": "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.\n\n2. The authors have mentioned limitations of their work anywhere.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are encouraged to hear that the reviewer found our work to be thorough and our method effective. Here are our responses to each point raised by the reviewer:\n- **On details about how the perplexity or BLEU scores were impacted by adding these optimizations.**\n\nOur optimizations will not negatively impact the perplexity or BLEU scores of the proposed model. Please check our common response 2 for a detailed explanation.\n\n- **On creating a separate limitation section.**\n\nDue to the space limit, we choose not to create a limitation section, but we have mentioned limitations of our method in the main text as well as the Appendix. Specifically, limitations of our method include:\n- Dynamic gating may suffer from the kernel launching cost, resulting in slightly higher latencies compared to Megablock when batch size is small. We discuss this on line 273-277, and a more detailed discussion can be found in App. D.1.\n- Expert buffering trades latency for a smaller memory footprint. Additional CPU-GPU communication is necessary for offloading the memory. We mitigate the effect by designing a buffering strategy based on the LLM activation pattern and applying load balancing. Discussion can be found at line 311-313.\n- We discussed other failed approaches we tried in Appendix Section E.\n\nWe hope that these responses help to clarify our points!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "n44QQa2lDz", "reviewer": "Reviewer_JXaG", "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.", "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": "- 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.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The paper provides a detailed characterization of MoE workloads, identifying specific sources of inefficiency.\n- The proposed techniques address key challenges in MoE inference.\n- The optimizations show substantial gains in throughput.", "weaknesses": "- 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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the very detailed review and suggestions. Here are our responses to each point raised by the reviewer:\n\n- **Including evaluation on representative configurations such as DeepSeek-MoE (2401.06066).**\n\nWe appreciate the reviewer’s suggestion to include a discussion on the DeepSeek-MoE architecture. We will add citations and a detailed discussion in a later version of this paper. It is difficult to have a fair comparison with DeepSeek-MoE, because (1) the DeepSeek-MoE released implementation assumes single-device deployment and does not support expert-parallelism, and (2) the DeepSeek-MoE released implementation is based on Huggingface TGI framework, whereas our work is based on NVIDIA Megatron, and the custom kernels involved made it hard for us to perform an apple to apple comparison.\n\nWe believe that our optimizations can also be applied to the DeepSeek-MoE architecture. Compared to traditional MoE Transformer architecture discussed in the paper, DeepSeek-MoE incorporates two additional strategies – (1) fine-grained expert separation and (2) the idea of shared experts. Below, we discuss the interaction of the two strategies and our optimizations.\n\n1. Our optimizations are relevant for DeepSeek-MoE because, even when the MoE layer activates multiple experts (e.g., six in 16B model), many of MoE’s inefficiencies remain. Expert sparsity remains a problem because the total number of experts is large, especially due to DeepSeek-MoE’s approach to fine-grained expert separation. In this setting, our optimizations for expert buffering and load balancing will reduce memory use and latency for multi-device inference. \n2. 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. \n\n- **On limitation of proposed optimization on small batch sizes compared to Megablock.**\n\nPlease check the common response 1 for a detailed response and discussion.\n- **On detailed results for memory and latency.**\n\nDetailed memory usage for each model and optimization combination can be found in Figure 11, Sec. D.2 of the Appendix. In our experiments, the sequence length for each task was fixed, making throughput effectively inverse of mean latency, thus we omitted a separate figure for it. We will provide a detailed table of the measured mean latency in future revisions.\n- **On readability and figure placing.**\n\nWe appreciate your understanding that the figures were moved to the appendix due to space constraints. We agree that the flow of the paper can be improved by relocating key figures back to the main text. We will prioritize moving more results to the main text as space permits in future versions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "p6JeQt9keO", "reviewer": "Reviewer_t1yX", "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.", "questions": "1. The authors show analysis for language modeling and machine translation. How about other tasks, e.g., summarization?", "limitations": "Limitations are sparsely discussed throughout the paper.", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The authors show consistent speed up over competing methods.\n2. The authors show detailed latency and memory analysis in their appendix.", "weaknesses": "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.\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the review and suggestions. The reviewer raised several questions which we address in order:\n- **Is this paper suitable for NeurIPS?**\n\nWe note that per the [call for papers](https://neurips.cc/Conferences/2024/CallForPapers), NeurIPS infrastructure track calls for submissions related to libraries, improved implementation and scalability, and distributed solutions, and we have submitted our paper to the infrastructure track. We believe our improved MoE implementation is a perfect match for this track.\n\n- **On performance impact.**\n\nDynamic gating and load balancing will have no negative impact on the accuracy, perplexity, or BLEU score of the model. See common response 2 for a more detailed explanation.\n\n- **On summarization performance.**\n\nWe note that summarization is increasingly performed by language models via zero-shot prompting (e.g. [1, 2]). Therefore, we believe experiments on summarization would be similar to those for language models.\n\n[1] Wu, Jeff, et al. \"Recursively summarizing books with human feedback.\" arXiv preprint arXiv:2109.10862 (2021).\n\n[2] Adams, Griffin, et al. \"From sparse to dense: GPT-4 summarization with chain of density prompting.\" arXiv preprint arXiv:2309.04269 (2023).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We would like to thank reviewers for providing us with valuable feedback. 

We have taken note of the concerns raised by each reviewer and addressed them in detail. Here, we provide responses to the most shared questions first followed by a detailed response to each reviewer's concern in the rebuttal.

- **Dynamic Gating performs worse than Megablock when the batch size is less than 32.** (BGdS, JXaG):

While Dynamic Gating performs worse than Megablock when the batch size is less than 32, we argue that such cases are actually rare during inference, especially online serving. Performing inference with small batchsize will limit the throughput due to lower computational intensity of the workload, which is demonstrated in Fig. 4. Therefore, to improve resource utilization, request serving will be separated into multiple steps, and during each step, requests will be dynamically **re-organized into large batches** (i.e., Continuous Batching), as proposed by [1]. This technique has been adopted by recent LLM serving systems, such as vLLM [2] and TGI [3]. In a production system, multiple requests arrive at the node with short intervals between their arrival times, creating numerous opportunities for continuous batching. Therefore, we argue that the **performance on larger batch sizes is more important** for online serving, and our technique outperforms Megablock in such a scenario. We also note that a detailed discussion on the reason for the performance trend is provided in Appendix. D.1. 

- **Impact of our optimizations towards the performance of the transformer model, in terms of perplexity, BLEU scores, or other metric of interest.** (dmgJ, JXaG, t1yX)

When expert capacity (the number of tokens accepted by an expert) is large, our optimized model is **mathematically equivalent** and as accurate as the baseline. When expert capacity is small, our optimized model may be more accurate (because it never drops tokens) than the baseline (which can drop tokens). Note that in baseline studies, evaluation is performed with large expert capacity and tokens are not dropped. 


[1] Yu, Gyeong-In, et al. "Orca: A distributed serving system for Transformer-Based generative models." 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22). 2022.

[2] Kwon, Woosuk, et al. "Efficient memory management for large language model serving with paged attention." Proceedings of the 29th Symposium on Operating Systems Principles. 2023.

[3] https://huggingface.co/docs/text-generation-inference/index

Meta-review: 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.
Decision: Accept (poster)