Scattered Mixture-of-Experts Implementation
Authors: Shawn Tan
Yikang Shen
Rameswar Panda
Aaron Courville
Venue: COLM (2024)
Keywords: triton, gpu, moe, mixture of experts, sparse mixture of experts
TL;DR: Triton-based implementation of Sparse Mixture of Experts without padded copies.

Abstract: ScatterMoE is an implementation of Sparse Mixture-of-Experts (SMoE) on GPUs. ScatterMoE builds upon techniques in existing implementations, and overcoming some of the current limitations to improve batched inference, training speed, and memory footprint. This implementation achieves this by avoiding padding and making excessive copies of the input. We also fuse expert linear transforms and reordering operations with ParallelLinear, a module that can be used to extend the concept of SMoEs. We benchmark our implementation  against Megablocks, and show that it enables a higher throughput and lower memory footprint. We also show how ParallelLinear enables extension of the Mixture-of-Experts concept by demonstrating with an implementation of Mixture-of-Attention.

Published as a conference paper at COLM 2024
Scattered Mixture-of-Experts Implementation
Shawn Tan
Yikang Shen
Rameswar Panda
aajings@mila.quebbs
Likaag.shenCibm.com
rpanda@ibm.com
Aaron Courville courvilaGiro ca
Abstract
ScatterMoE is an implementation Sparse Mixture-of-Experts (SMoE) on GPUs. ScatterMoE builds upon techniques in existing implementations, and overcoming some of the current limitations to improve batched infer- ence, training speed, and memory footprint. This implementation achieves this by avoiding both padding and making excessive copies of the input. We also fuse expert linear transforms and reordering operations with Par- allelLinear, a module that can be used to extend the concept of SMoEs. We benchmark our implementation against Megablocks, and show that it enables higher throughput and lower memory footprint. We also show how ParallelLinear enables extensions of the Mixture-of-EEppert concept via a demonstration with Mixture of Attention implementation.
http:///gthha.coom///aattttttttt
1 Introduction
Sparse Mixture of Experts (SMoEs; Shazeer et al. 2017) have become increasingly popular for scaling up Transformer-based language models. While applications of SMoEs like the Switch Transformer (Fedus et al., 2022) use SMoEs to scale "outrageously" large models by distributing computation of experts across compute nodes, it has proven useful even in scaling up smaller models where device memory is an issue.
For SMoEs, sparsity key in reducing computation costs. However, fully exploiting the sparsity to improve the throughput of MoË modules is challenging. While a lot of deep learning research implemented in PyTorch (Paszke al., 2019), the naive implementation of SMoEs does not take full ad advantage of the parallelism of GPUs and are slow result. Initial implementations on Tensor Processing Units (TPUs) for Switch Transformers require all tensor sizes, capacity, to be specified at compilation, which ensures that all load for every expert is equal (Fedus et al., 2022). This creates issues when experts are imbalanced: When the router assigns more tokens than capacity allows to particular expert, some tokens are dropped. Likewise, when experts are underused, the tensors are padded, which creates unnecessary memory allocation. Later, Megablocks (Gale et al., 2023) and PIT (Zheng et al., 2023) framed the SMoE computation as sparse matrix multiplication problem, which can be computed efficiently with sparse matrix optimised algorithms. In both these cases, the authors were able to create more efficient GPU-based implementation of SMoEs.
Despite these recent advances, there still room for improvement. First, existing implemen- tations of SMoEs, performs a scatter-to- to-orrtto-rgupppp initial copy of the input, creating memory allocation overhead during training because of tensors stored for used in the backward pass. Some implementations pad the routed tensors so they are equal-sized blocks, which further increases the memory overhead. Second, Megablocks and PIT requires a translation of the SMoE problem into sparse matrix format. While this incurs only a small part of the computation overhead, the sparse matrix format makes the intermediate representation harder to extend upon.

---

Published as a conference paper at COLM 2024
Copy Group
Linear transform
Non-linearily
3
Linear transform
2
Copy Scaller
0
2
33
Other Implementations
ScatterMoE
Figure 1: Current implementaionn of SMoE Multi-layer Perceptrons (MLPs) require copy of the embeddings when grouping (left), while ScatterMoE fuses the grouping and linear transformation step (right), reducing the memory footprint of our method. The various colours represent different experts, while the vertical rectangular boxes represent embeddings with their associated time steps labelled above or below them.
In this paper, we present ScatterMoE, an SMoE implementation that minimises this memory overhead. This made possible by ParallelLinear, a primitive we introduce that performs grouped matrix operations on scattered vectors. The resulting intermediate representations (eg. hidden state of an SMoE MI.P) can be exposed as standard PyTorch tensors, allowing for easy extensions of present SMoE methods to other types of expert modules. We demonstrate the utility this representation by implementing SMoE Attention with ParallelLinear, following the specification in Tan et al. (2023). In the final section, we benchmark ScatterMoB against a naive PyTocch implementation and Megablocks.
Related Work
Other implementations & Dependencies The core parts of ScatterMoE. is implemented with Triton' (Tillet et al., 2019), tile-based language for GPU programming in Python, making it the most accessible for modification and extension. Our main comparison is against Megablocks?, which is implemented using the STK framework3 which also uses Triton. Megablocks is also used in the Megatron-LMM model (Shoeybi et al., 2019; Narayanan et al., 2021; Korthikanti et al., 2023), and its widespread use as efficient method for training SMoEs. Another popular library for implementing SMoEs is CUTLASS et al., 2022), with which Megablocks uses as its grouped option.
httpe///gttuu..co//oeeeas/tttatt
http://gtthab.comm//VDDDD/aaua
2

---

Published as a conference paper at COLM 2024
3
(a) Grouped to grouped
(b) Scattered grouped
2
3
1 2
2
2
(c) Scattered to scattered
(d) Grouped to scatter
Figure 2: ParallelLinear allows for performing different combinations of SMoE transfor- mations allowing for the input and output to be either grouped or scattered. This basic functionality forms the basis of both forward and backward passes of ScatterMoE. Unlike existing implementations, these operations are performed without additional copying (or padding) of the input and output tensors.
Other applications of SMoEs Aside from MLPs,, SMoE versions of the attention module have also been proposed (Zhang et al., 2022; Csordás et al., 2023). These mixture-of-attente (MoA) implementations have been used to scale up Universal Transformers (Dehghani et al., 2018; Tan et al., 2023), and also for applications to continual learning in fully modularised Transformer (Shen et al., 2023). Computing SMoEs efficiently will bring huge benefits to the training and inference of these models.
Method
In this paper, will maintain the notation convention of boldface for matrices, i.e. X. Unless otherwise stated, the first dimension is batch-time (batch and time dimensions flattened) for ease of understanding. Additionally, X; denotes the i-th row of X.
3.1 Sparse Mixture-of-Experts
SMoE modules made up of E experts which typically sub-modules of similar architecture. Each of the T tokens in the input is routed via a routing module, and then based on the router output weights, assigned method of computing an SMoE (iterating over all tokens and evaluating the respective expert output) far slow, and does not exploit the full parallelism of GPU computation. In practice, SMoE implementations often perform the following main steps:
experts, where k < E. However, the naive
1. Routing - Based on each token embedding X,, the router assigns the weights for each expert g (X,), and only the top-k experts are selected.
2. Grouping- This step groups all tokens that are assigned to the same expert together. I 1, as is often the case, then this also results in "fan out" of tokens, resulting in kN embeddings in total.
3. Expert transform Now that the tokens are grouped by expert, each expert (a linear transform) can be efficiently computed by batched vector transformations (matrix-matrix multiplications).
4. Scattering This step returns each token to be grouped by its original time-step. This still results in kN embeddings in total.
3

---

Published as a conference paper at COLM 2024
5. Weighted sum This step combines the k outputs per token by its original routing weight,
Y
E
ge (Xt) fe (Xt)
ectopk(g(X)))
resulting in N embeddings.
Typically, each expert is a Multi-layer Perceptron (MLP) with one hidden layer dimension dexpert.
In Megablocks, steps (1) and (4) result in a copy of the original input. They further pad the per-expert blocks of tokens so that they fit into equal count for convenient GPU computation (See Figure 1). This allocates the padded array for the embeddings in High Bandwidth Memory (HBM) and copies the original embeddings in sorted order into it. A Grouped GeMM then performed on the expert-sorted array.
ScatterMoE, the other hand, avoids realising the entire padded array in HBM. Instead of copying all the embeddings into a padded array, we sort the tokens according to the experts, and pad the indices instead. When loading a tile into Static RAM (SRAM), load according to the padded indices, resulting in padded tile.
3.2 ParallelLinear operation
Our implementation of SMoE relies on ParallelLinear, which allows for different combina- tions of grouped General Matrix Multiplications (GeMMs). In order to achieve this, wrote Triton kemnel, scatter2scatter, that enables all combinations of operations shown in Figure This operation fuses grouped GeMMs and scattered read and write operations, which allows to skip an intermediate group and copy step. ParallelLinear allows options for grouped and scattered for both input and output, resulting in the four possible combina- tions seein in Figure 2. With combinations of these operations, we can implement both the forward and backward passes of ParallelLinear.
Algorithm provides the pseudo-code of ParallelLinear. It is a thin wrapper around the scatter2scatter kernel, and the grouping options are provided as arguments to Paral- lelLinear. This is the main workhorse of our SMoE implementation, and can be used for both implementing an MLP and an attention layer. We implemented the backward pass of ParallelLinear independently of its downstream usage to allow it to be used primitive to build other experts upon.
A Algorithm 1 ParallelLinear FORWARD
Input: X
T din
input matrix transform tensor
order indices top-k
E x din x dout
k
default = 1
Sxj
routing weights default p: (Ikx
where Sj Tk
*options
grouped
True, False
grouped.out
True, False
Output:
Y S dout
output matrix
Y scatter2scatter (X, o,k *options)
p # then
Y view SY, -1)
reshape and weighted sum if is provided.
Y bmm (p. Y)
else
Y
end if
4

---

Published as a conference paper at COLM 2024
3.2.1 Backward pass
In a typical batched linear transformation XW, need compute the gradients VX and VW.
VX= VYWT,
VW X VY,
In ParallelLinear, we will need compute these gradients for each of the E experts. While this could be computed in way where X and V are both scattered, the implementation of this operation is fastest when both X and VY are grouped
Algorithm& 2 ParallelLinear BACKWARD
Input:
VY
S x dout
gradient matrix
Tk dout T x din Tk
from forward pass from forward pass from forward pass
W
E x din x dout
transform tensor
k
top-k
defiultt 1 routing weights default p x 1)==
column p
S
where Sj =
Output:
VX
same size as X same size p
VW
same size as W
Vp
if p # Ø then
Vp bmm VY, Y)
VY group VY,
I| weight and group
else
VY - Y
end if
if X is not grouped then
X group(X.o,0)
else
X 4- x
end if
VW groupXTY (X, VY) VY, W 0.
VX + scatter2scatter
grouped to scaffer or group depending on original input
Algorithm 2 groups the embeddings when they not grouped in order to compute VW. groupXTY is the kernel that implements this grouped matrix multiplication. While this grouping incurs additional memory allocations for potentially very large matrices, these array allocations could be reused. Once the gradients Vp has been computed, the array for can be used as the output for VY. X can be reused for VX as they are of the same dimensions. We can then further minimise the use of memory during the backward pass by re-using the arrays used for the grouping operations. We colour the reused arrays in blue and orange respectively Algorithm 2.
3.2.2 SMoE Multi-layer Perceptron (SMoE MLP)
In the context of an SMoE MLP, we reduce the memory footprint even further. The MLP requires two linear transformations, and could be naively implemented with two ParallelLinear operations set perform scatter-to-scater transformations, However, we can configure these two linear transforms to be scattered-to-grone scattered respectively. This will allow each ParallelLinear transform in the SMoE MLP to require only one group operation during the backward pass.
then grouped-to-
3.3 Extensibility: Mixture-of-Attenting (MoA)
There have been several proposals for applying SMoEs to the attention layer (Zhang et al., 2022; Csordás et al., 2023; Tan et al., 2023). Regardless of formulation, before the attention
We tested scatterXTY and found it be slower than grouping operation followed by groupXTY

---

Published as a conference paper at COLM 2024
Algorithm 3 SMOE MULTI-LAYER PERCEPTRON
Input:
x dmodel
input matrix transformation tensor transformation tensor
o
vector grouped ordering
W1
E x dmodel x dexpert
W2
Ex x dexpert model
Output:
T dmodel
output matrix
H + ParallelLinear (X, W1, 3 grouped Falsin grrpaalle.guuud..ourro outee
H <- 0 (H)
// where ø is any point-wise non-linearity
Y ParallelLinear (H, W2, o, p. p.grouped -rrouped_i--Truee.e hrapppddddnnnu..ooopddoooo---aa grouped
layer is applied, the embeddings should be in chronological order (scattered) to facilitate the application of positional embeddings and to compute the result of the attention weights and value embeddings. With existing SMoE implementations, there would be additional pair of group-scatter operations, incurring additional memory costs.
ScatterMoF provides ing through a ParallelLinear transform, we can implement MoAs without allocat- ing the extra arrays for grouping and scattering. Figure 3 shows the opera- tions used for SMoE Attention. In this report, we specifically implement and benchmark the Mixture of Multi-head Attention variant found in Tan et al. (2023). Attention (GQA;
advantage. Since we can retain the scattered order-
This version resembles Grouped-ques
Ainslie et al. 2023), where each key head has multiple possible query heads, while in the SMoE setting, there would be Nexxeert key heads, with k hexpert query heads selected from a possible E hexpert heads, where hexpert is the number of heads per expert.
Following the standard attention implementation con- ventions, we set dout = dexpert dhead, where dhead is the number of dimensions for each head. We then reshape accordingly when we perform the attention operation so that the separate heads interact independently.
Muti-head
In Algorithm 4, we also consider the time dimension, so we express the inputs and intermediate arrays 3- dimensional tensors with the dimensions for batch, time, and embedding dimensions (B T x d). In practice we assume that the input is contiguous and is batch-time ordered, allowing us to flatten the tensor and proceed as we did in the case of the ML.. Note that for the SMoE attention, a key distinction is that we require ParallelLinear to give an ungrouped output after the first transformation, and A
Figure 3: ParallelLinear allows for scattered to scattered trans- formations which retains the chronological order.
it takes ungrouped input for the output transform, which means both ParallelLinear transformation use a scattered to scattered configuration (Figure 2c).
4 Performance
In this section, we cover the performance of our implementation for both training and inference. As an overall integrated test, we benchmark our method within Mixtral (Jiang et al. 2024), using 1.5B parameter configuration, dmodel 1024,
dexpert 3584, k E=8,
L = ==16,
We compare against the naive implementation from HuggingFace (Naive HF impl.), then swapping out the SMoE layer with Megablocks sparse (MB (Sparse)) and grouped memory efficient (MB (Mem. eff.)), and finally ScatterMoE (Ours). Our goal is to measure the overall throughput in training setting.
We ran the training for 100 training updates, with an effective batch size of 256 and 2048 tokens per instance, across 8 A100 GPUs the same compute node. For both the naive and

---

Published as a conference paper at COLM 2024
Algorithm 4 SMOE MULTI-HEAD ATTENTION
Input:
X
x BxTx mmodel BTK
input matrix
k
top-k router weights value transform output transform
grouped indices key transform
X
W
moddel x dout
W V Wo
dmodel X dout E dout X model
W Q
E x model dout
query transform
Output:
B x x model
output matrix
C Wv
W
ParallelLinear (X, W. o, www.pppeiiiaaaaeeeerrreaee..o--a groppeeo.oooooee..oot--a.a-e
t Attention (Q, K, V)
ParallelLinear (Ó, Wo, o, p. grouped False,, grtt
proupp out Fut=False]
1e6
Jeg
Throughpui (takens/5)
Naive (HF impl.)
Ours
Toouge ou bakena
1.25
Megablocks PyTorrh
Megablocks
200000
MB (Sparse)
1.00
PyTorch
Ours
0.75
Meeorry
100000
0.50
0.25
0.00
Backeward
Forward Backward
(a) 1.5B model training throughput (b) SMoE MLP unit throughput (c) SMoE MLP unit memory use
Figure
ScatterMoE, this resulted in an actual batch size of 128 and 2 accumulation steps, while the Megablocks benchmarks required a batch size if 64 and 4 accumulation steps. We ran the training for 100 steps and computed the overall throughput. Our method outperforms both the Sparse Megablocks implementation by 38.1% % in this setting. This indicates the importance of the smaller memory footprint, but at larger dimensions with equivalent batch sizes, the gains are not as significant.
The rest of this section covers the other benchmarking experiments we performed in more detail, testing for the effects of decreasing sparsity, and increasing granularity of experts, and benchmarking of our Mixture of Attention implementation.
4.1 Unit Benchmarking on the SMoE MLP
Unless otherwise stated, we the following model hyperparameters,
dmodel 409,,
dy model
dexpert d4/k,
E 8k
For example, the active hidden units for MLP is - 4096 8192. If 4, then E 4
32, with each expert being 8192/k = 2048 dimensions. Each datapoint on the plot is the median and 5-th and 95-th percentiles of 100 runs of the module. In this unit test, we use more efficient PyTorch implementation from the implementation of Tan et al. (2023)
Figure 4a summarises the overall performance for an SMoE MLP where E = 32, k 4, and T= 30 2048. All benchmark times are measured on an Nvidia A100 GPU. Generally, we find that ScatterMoo has slightly higher throughput during training, for the same input sizes. Our method shows a larger margin of improvement for inference.
On memory consumption, implementation for the SMoE MLP uses 66.2% of the memory Megablocks uses during training, while using only 53.6% the memory of Megablocks consider only inference.
httt:////tttuccc///////////////e layer/parallel linear/parallel experts

---

Published as a conference paper at COLM 2024
75
Megablacks
75
Megablacks
70
Relative % throughout
70
throughput
65
60
60
55
50-
2
10 12 14 16
8 10 12 14
16
G
(a) Training throughput
(b) Inference throughput
Figure Increasing k and E while fixing the number of active parameters and total param- eters. We find that our implementation scales better with higher k. Inference granularity scaling performance. The difference in relative throughput higher we consider only the forward pass.
4.2 Granularity and Throughput
Krajewski et al. (2024) defines the concept of granularity. Given a SMoE MLP with the equivalent active parameters MLP with intermediate dimension layer of dimension dg, and with each expert of dimension despert, then granularity is defined as,
dg dexpert
Here, measure the effect of throughput we vary G while fixing dg. Accordingly, with higher values of G, we need to increase values of k and E more granularity requires more experts to achieve the same active parameters.
We test values of k E [1,2,4,8 16), and E = 8k for divisibility of dimension sizes. Figure 5 shows how throughput of both methods vary relative to model with equivalent active parameters, Since we maintain constant active and total parameters for these runs, these results are also constant for all G.
We find that ScatterMoE scales with G with better throughput. This may be related to the increase in zero-padding required for higher number of E in the case of Megablocks as the number of experts grows and the embeddings assigned to each expert decreases, there will be more padding introduced. If we consider only the forward pass, the relative gap between our method and Megablocks is also much higher than in the case of training. This makes our method favourable for batched inference, especially with high granularity settings. The results of Krajewski et al. (2024) suggests higher G for SMoE models, and our implementation seems suited for this application.
4.3 Decreasing sparsity
We can view the SMoE as an interpolation between a model with the size of just the active parameters k- dexpert and large fully dense model with dexpert. However, SMoE comes with additional overhead (e.g. routing, sorting), and we want to measure how much fully dense model.
reducing sparsity will affect throughput, in comparison
In this experiment, we tested on increasing values of < 30. Further increasing reaches the limit device memory for Megablocks. We maintain E 64 for all runs, and compare the performance of both Megablocks and ScatterMoB to dense model with d = E dexpert-

---

Published as a conference paper at COLM 2024
1e6
Throughius (tolensis)
Ours
Throughpat (takertit))
450000
Throuphoput [tokena/S]
1.3
Megablocks
1.00
400000
0.75
1.1
0.50
350000
0.25
Ours
0.9
Ours
Megablocks
Megablocks
300000
Forward
Backward
Both
(a) MoA throughput
(b) Training throughput
(c) Inference throughput
Figure 8: The curves for increasing granularity for the MoMHA implementation. In this case, the Active params baseline varies in throughput because of the shared key and value vectors across the experts.
relative % throughput
1500
Same params.
Generally, we find that while implemen- tation performs with slightly better through- put overall, both Megablocks and our imple- mentations are still more efficient than large dense MLP with equivalent parameters. How- ever, note that in this case, the throughput for k 30, 64 is already reaching the through- put for an equivalent dense model with the same parameters. Further increasing exceeds the memory of the device ran the benchmarks on for Megablocks.
Ours
Megablocks
1000
500
10
20
30
k
Figure 6: Relative throughput curves as we decrease sparsity (increasing k)
4.4 Mixture-of-Attention
As previously mentioned, we implement the Mixture Multi-head Attention (MoMHA) variant of Attention SMoEs as described in Tan et al. (2023). This implementation allows for multiple attention heads per expert, and shares the key and value embeddings across the attention experts. This formulation is similar to Grouped-query -uery Attention (GQA; Ainslie et al. 2023), where each head the key has several query heads, each these forming group. MoMHA experts the equivalent of groups in GQA, where each group is of size k.
For the following benchmarks, we adhere to the following parameters:
dmodel 4096, dhead 128, T 2048, h 32, hexpert h/k,
E 8k
where dhead is the dimension of each attention head and the number of active attention heads.
We implemented an equivalent baseline in Megablocks using the 'dense' configuration in the library. This version still suffers from the issue with having to perform redundant grouping and scattering steps.
For = 8, we find that our implementation outperforms Megablocks, by 24.0% of througput for inference. We also note from Figure 8, that as we increase granularity (fewer heads per expert smaller hexpert), the gap between implementation and Megablocks grows as well. Again, our method is favourable for use high granularity settings for Mixture-of- Attention

---

Published as a conference paper at COLM 2024
Tasks
Metric
Hugging Face
ScatterMoE
Abs. Error
winogrande
0.7632
0.7640
0.0008
sciq
0.9520
0.9580
0.0060
race
0.4057
0,4010
0.0047
piqa
0,8330
0,8368
0.0038
openbookqa
0.4680
0,4740
0.0060
hellaswag
accuracy
0.8396
0,8405
0.0009
copa
0,9300
0,9300
0.0000
boolq
0.8523
0.8541
0.0018
arc_easy
0.8350
0.8350
0.0000
arc.challenge
0.5973
0.5981
0.0008
wikitext perplexity
5.6135
5.6142
0.0007
Table 1: Language Model Evaluation Harness results comparisons: Hugging Face v. Scatter- MoE implementations. Differences in results between both implementations negligible.
4.5 Mixtral Inference Comparison
Finally, converted Mixtral 8x7B² to ScatterMoE, and ran the LM Evaluation Harness (Gao et al., 2023) on several benchmarks (See Table 1). As ScatterMoE is an alternative implementation of Sparse MoEs, do not expect any differences in the final evaluation results. The absolute error between the Hugging Face naive implementation and ScatterMoB is sufficiently small, which demonstrates this.
We have included the script to convert the model parameters from Hugging Face's format to a format compatible for ScatterMoE
5 Conclusion & Limitations
ScatterMoF, is an implementation of SMoEs in Triton that reduces the memory footprint and offers slightly higher throughput on the GPU compared to existing solutions. We have also engineered ScatterMoE to use the ParallelLinear primitive, which we envision to be module that can extended upon to build other SMoE-style modules that require grouped scattered linear transformations.
At present, ScatterMoE does not implement a specialised kernel for speeding up decoding, and further work is required for parallelisation in multi-node setting. These additional features will be added future iterations, and we believe further testing by us and the open source community will iron out any remaining bugs and performance issues left unoptimised.
Finally, we have also provided an implementation of an MLP and Attention layer based on ScatterMoE, that we hope will benefit any future implementations of Mixture-of-Exper based models, and serve worked examples for extending the concept of SMoEs to other variants of linear transformation based experts.
Acknowledgments
We would like to Bowen Pan and Mayank Mishra for their testing and feedback of Scatter- MoE. Songlin Yang also gave valuable advice during the development of the ScatterMoF kernels.
7 /nttpss//nuggingfaca.. httt:///wggaaaaaeecc//mm/aaaaala
10

---

Published as a conference paper at COLM 2024
References
Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arviv preprint ar Xrio:2305.13245 2023.
Róbert Csordás, Piotr Piekos, and Kazuki Trie. Switchhead: Accelerating transformers with mixture-of-experts attention. ariiv preprint arXio:2312.07987, 2023.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. arXiv preprint ar arXiv:1807.0388 2018.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. The Journal of Machine Learning Research, 23(1):5232-522) 5270, 2022.
Trevor Gale, Deepak Narayanan, Cliff Young, and Matei Zaharia. MegaBlocks: Efficient Sparse Training with Mixture-of-Expers.. Proceedings of Machine Learning and Systems, 5, 2023.
Leo Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noachh, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 12 2023. URL
Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprint ar Kio:2401.04088, 2024.
Young Jin Kim, Rawn Henry, Raffy Fahim, and Hany Hassan Awadalla. Who says elephants can't run: Bringing large scale moe models into cloud scale production. arXiv preprint arXiv:2211.10017, 2022.
Vijay Anand Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, and Bryan Catanzaro. Reducing activation recomputation in large transformer models. Proceedings of Machine Learning and Systems, 5, 2023.
Jakub Krajewski, Jan Ludziejewski, Kamil Adamczewski, Maciej Pióro, Michat Krutul, Szymon Antoniak, Kamil Ciebiera, Krystian Król, Tomasz Odrzygóżdł, Piotr Sankowski, et al. Scaling laws for fine-grained mixture of experts. ar preprint ar arXiv:2402.0700 2024.
Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Pat- wary, Vijay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, et al. Efficient large-scale language model training gpu clusters using megatron-Im. In Proceedings of the International Conference for High Performance Computing, Nefworking, Storage and Analysis, pp. 1-15, 2021.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An im- perative style, nigh-performance deep learning library. Advances in neural information processing systems, 32, 2019.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gate mixture- of-experts layer. ar preprint ar Xio::7011.05534 2017.
Yikang Shen, Zheyu Zhang, Tianyou Cao, Shawn Tan, Zhenfang Chen, and Chuang Gan. Moduleformer: Learning modular large language models from uncurated data. Xiv preprint ar (tte2330004640) 2023.
11

---

Published as a conference paper at COLM 2024
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-lm: Training multi-billion parameter language models using model parallelism. arvio preprint ar Stt::1999000055 2019.
Shawn Tan, Yikang Shen, Zhenfang Chen, Aaron Courville, and Chuang Gan. Sparse universal transformer. ar arv preprint ar Xrix::2310.00066 2023.
Philippe Tillet, Hsiang- Tuung Kung, and David Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, pp. 10-19, 2019.
Xiaofeng Zhang, Yikang Shen, Zeyu Huang, Jie Zhou, Wenge Rong, and Zhang Xiong. Mixture of attention heads: Selecting attention heads per token. arXiv preprint arXiv:2210.05144, 2022.
Ningxin Zheng, Huiqiang Jiang, Quanlu Zhang, Zhenhua Han, Lingxiao Ma, Yuqing Yang, Fan Yang, Chengruidong Zhang, Lili Qiu, Mao Yang, et al. Pit: Optimization of dynamic sparse deep learning models via permutation invariant transformation. In Proceedings of the 29th Symposium on Operating Systems Principles, pp. 331-347, 2023.
12


Reviews:
[{"review_id": "LdRPWInH9A", "reviewer": "Reviewer_9gxi", "summary": "This paper proposes ScatterMoE which optimizes the memory cost and the throughput of tokens for SMoE models. Specially, ScatterMoE focuses on the issue on extra memory copies of token representations during the operation of MoE scattering and gathering. By implementing custom kernels using Triton or CUTLASS, ScatterMoE eliminates the extra memory copies and improves the overall training & inference performance, compared with existing SMoE frameworks like Megablocks. The work provides the detailed implementation of ScatterMoE and extensively compares the performance with baselines (Megablocks and naive PyTorch implementation) considering a variety of experimental settings (granularity, sparsity and MoA).", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback!\n\n> As a competitor of Megablocks,... , I think the work will be much more influential and have more significance to be accepted.\n\nWe agree with this perspective,  and ScatterMoE has been integrated into an open source fine-tuning framework and we have plans to include it into a pretraining framework are underway (pull-request submitted). \n\nHowever,  ScatterMoE is a modular, light-weight implementation of MoEs that is meant to  be used as a drop-in replacement for MoEs. We’ve added a conversion script for the Mixtral model on huggingface to demonstrate this, and include the lm-evaluation-harness results achieved with both the naive huggingface implementation and using ScatterMoE, and the results show that the differences are negligible. \n\n|     Tasks     |      Metric     |   HF   | Scattermoe |\n|:-------------:|:---------------:|:------:|-----------:|\n| winogrande    | acc             | 0.7632 |     0.7640 |\n| wikitext      | word_perplexity | 5.6135 |     5.6142 |\n| sciq          | acc_norm        | 0.9520 |     0.9580 |\n| race          | acc             | 0.4057 |     0.4010 |\n| piqa          | acc_norm        | 0.8330 |     0.8368 |\n| openbookqa    | acc_norm        | 0.4680 |     0.4740 |\n| hellaswag     | acc_norm        | 0.8396 |     0.8405 |\n| copa          | acc             | 0.9300 |     0.9300 |\n| boolq         | acc             | 0.8523 |     0.8541 |\n| arc_easy      | acc_norm        | 0.8350 |     0.8350 |\n| arc_challenge | acc_norm        | 0.5973 |     0.5981 |\n\n> It will be much better if larger scales of MoE models are benchmarked....\n\nWe benchmarked a bigger model of 9B parameters: 32 experts and k=4, d_model = 2048, mlp_size = 1536 per expert, 26 layers. This was the largest we could manage with a single node and 8 X A100 80gb GPUs. At this size, the throughput for a minibatch of 32 * 2048, with 16 acc steps for an actual batch size of 2, was 57957.8 tokens/s for ScatterMoE, and 19924.9 tokens/s for Megablocks scattered. However, this may be due to a memory leak in the recent version of Megablocks. \n\n> For some recent variants of MoE like shared experts (DeepSeekMoE for e.g.),...\n\nScatterMoE can support the shared expert scheme in DeepSeekMoE: one simple way to do this is to either adjust the logit to a large value for that expert, or to append the shared expert index before passing it to the MoE module.  However, it is unlikely that there will be an advantage over those already discussed in this report.", "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": "1. The motivation is straightforward and reasonable, which directly eliminates the excessive memory consumption of representation tensor copies and token paddings.\n2. The paper describes the implementation clearly including the forward & backward passes of ParallelLinear and its extension to MoA modules.\n3. The comparison in performance provided is extensive and includes sufficient conditions (like granularity and sparsity) which affects the experimental results.", "reasons_to_reject": "1. As a competitor of Megablocks, this work should better be accomplished into a usable package (instead of just raw fused kernels) and provide a more end-to-end time-basis training curve for current prevalent LLMs. If achieved, I think the work will be much more influential and have more significance to be accepted.\n2. It will be much better if larger scales of MoE models are benchmarked, in fact a 1.5B parameter configuration is more like a toy setting.", "questions_to_authors": "For some recent variants of MoE like shared experts (DeepSeekMoE for e.g.), will ScatterMoE still be supported and in advantage in performance?", "extra_scores": {}}, {"review_id": "LascJPrx5M", "reviewer": "Reviewer_4BFX", "summary": "This paper introduces a new method called ScatterMoE for the efficient implementation of MoE. It proposes a new kernel called scatter2scatter which supports group-to-group, group-to-scatter, scatter-to-scatter, and scatter-to-group operation. The benchmark against the recent SOTA work MegaBlocks demonstrates its higher throughput and lower memory footprint.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the feedback!\n> The paper should provide more information about the details of the key innovation, including its detailed implementation, how it avoids padding, instead of using paragraphs to introduce MoE, its integration with MoE and MoA. Also, the pseudo-codes are not informative enough.\n\nWe will include more detailed information about the specifics you pointed out. Specifically, we will expand on the following points in the paper:\n1. Instead of performing a padded copy in High Bandwidth Memory (HBM), we sort the tokens according to the experts, and pad the indices instead.\n2. When loading a tile into Static RAM (SRAM), we load according to the padded indices, resulting in a padded tile. The full padded array of token embeddings is never realised in HBM.\n3. We will add the pseudo-code for the scatter2scatter kernel itself, since the current pseudo-code only describes the overarching process using scatter2scatter as a primitive.\n\n> It is advised to provide more experimental results about model quality to ensure that using ScatterMoE can achieve similar or even better quality on the LLM evaluation benchmarks.\n\nWe converted Mixtral 7B to use ScatterMoE (script included), and ran the lm-evaluation harness on the following benchmarks. \n\n|     Tasks     |      Metric     |   HF   | Scattermoe |\n|:-------------:|:---------------:|:------:|-----------:|\n| winogrande    | acc             | 0.7632 |     0.7640 |\n| wikitext      | word_perplexity | 5.6135 |     5.6142 |\n| sciq          | acc_norm        | 0.9520 |     0.9580 |\n| race          | acc             | 0.4057 |     0.4010 |\n| piqa          | acc_norm        | 0.8330 |     0.8368 |\n| openbookqa    | acc_norm        | 0.4680 |     0.4740 |\n| hellaswag     | acc_norm        | 0.8396 |     0.8405 |\n| copa          | acc             | 0.9300 |     0.9300 |\n| boolq         | acc             | 0.8523 |     0.8541 |\n| arc_easy      | acc_norm        | 0.8350 |     0.8350 |\n| arc_challenge | acc_norm        | 0.5973 |     0.5981 |\n\nThe difference in the results are negligible, which we expect as ScatterMoE is a different implementation and does not make any fundamental changes to MoEs.", "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": "1. From the interpretation in the paper, this method is simple by fusing fuses grouped GeMMs and scattered read and write operations and effective by higher throughput and lower memory footprint.\n2. The paper is clearly written and easy to read.", "reasons_to_reject": "1. The paper should provide more information about the details of the key innovation, including its detailed implementation, how it avoids padding, instead of using paragraphs to introduce MoE, its integration with MoE and MoA. Also, the pseudo-codes are not informative enough.\n2. It is advised to provide more experimental results about model quality to ensure that using ScatterMoE can achieve similar or even better quality on the LLM evaluation benchmarks.", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "A4bkAahfOF", "reviewer": "Reviewer_dRgj", "summary": "This paper presents ScatterMoE, an SMoE implementation that minimizes memory overhead. Since ParallelLinear allows for performing different combinations of SMoE transformations allowing for the input and output to be either grouped or scattered, ScatterMoE relies on ParallelLinear to achieve both forward and backward passes. Generally, these operations are performed without additional copying (or padding) of the input and output tensors involved in the existing implementations. Experiments validate the effectiveness of ScatterMoE on throughput and memory use.", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> For Megablocks and PIT, the authors stated that “the sparse matrix format makes the intermediate representation harder to extend upon”. Does this lead to memory overhead? If yes, can you describe the reasons? Otherwise, is there any in-depth analysis or experiments on the aforementioned feature of Megablocks and PIT?\n\nIn an MoE MLP, the sparse matrix format does not create the overhead.  Megablocks makes a re-ordered copy of the input before performing the GeMM operation, which is the step that creates the additional memory overhead.  Figure 1 in the paper illustrates the additional copying that causes this.\n\nHowever, extending MoE to Attention the memory issue is further compounded:\n\n1. First we need to perform a ParallelLinear transformation of the input, to produce queries, which will require re-ordering to *expert order.* \n2. These queries will need to be re-ordered to *chronological order* for application of position embeddings (e.g. RoPE) and appropriate masking.\n3. Once the standard attention step is done, the output embeddings will have to be reordered to *expert order* again, \n4. before the final output expert transformation, before being reordered to *chronological order*. \n\nThere are 4 re-ordering steps in total, and each will incur an additional duplication of the input. With scatter2scatter, we can bypass the reordering by only performing the expert ordered operation in SRAM. This results in a 42% savings in memory usage.\n\nBoth Megablocks (in sparse mode) and PIT rely on a sparse representation of the intermediate hidden state in an MoE MLP. Unless the operation performed on the sparse representation is a point-wise operation (like an activation function, for example) to perform anything that requires order sensitivity (attention) would require converting back to the dense format, then re-ordering, incurring all of the above-mentioned costs AND the sparse-dense-sparse conversion. \n\n> In the comparative experiments, what is the performance of PIT?\n\nWe were unable to find an open source implementation of PIT to compare against", "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": "1) The implementation proposed in this study is beneficial for various real-world applications.\n\n2) The paper is written-well with a good presentation style and experimental results.", "reasons_to_reject": "I have several questions to authors, as listed in the following part.", "questions_to_authors": "1) For Megablocks and PIT, the authors stated that “the sparse matrix format makes the intermediate representation harder to extend upon”. Does this lead to memory overhead? If yes, can you describe the reasons? Otherwise, is there any in-depth analysis or experiments on the aforementioned feature of Megablocks and PIT?\n\n2) In the comparative experiments, what is the performance of PIT?", "extra_scores": {}}]

Author rebuttal:

Meta-review: This work proposes a new approach for a more efficient implementation of sparse MoE models by reducing memory copies, particularly compared with existing frameworks such as Megablocks.

Pros: The approach is straightforward and essentially works by removing extraneous memory copies. The approach is cleanly described and there are thorough performance comparisons with alternative approaches such as Megablocks. The approach performs well even up to 9B model size.

Cons: It would be great to have this integrated into a useable pretraining package.

Overall, I think this would be a great contribution to the conference.
Decision: Accept