TA-MoE: Topology-Aware Large Scale Mixture-of-Expert Training
Authors: Chang Chen
Min Li
Zhihua Wu
Dianhai Yu
Chao Yang
Venue: NeurIPS 2022 Accept (2022)

Abstract: Sparsely gated Mixture-of-Expert (MoE) has demonstrated its effectiveness in scaling up deep neural networks to an extreme scale. Despite that numerous efforts have been made to improve the performance of MoE from the model design or system optimization perspective, existing MoE dispatch patterns are still not able to fully exploit the underlying heterogeneous network environments. In this paper, we propose TA-MoE, a topology-aware routing strategy for large-scale MoE trainging, from a model-system co-design perspective, which can dynamically adjust the MoE dispatch pattern according to the network topology. Based on communication modeling, we abstract the dispatch problem into an optimization objective and obtain the approximate dispatch pattern under different topologies. On top of that, we design a topology-aware auxiliary loss, which can adaptively route the data to fit in the underlying topology without sacrificing the model accuracy. Experiments show that TA-MoE can substantially outperform its counterparts on various hardware and model configurations, with roughly 1.01x-1.61x, 1.01x-4.77x, 1.25x-1.54x improvements over the popular DeepSpeed-MoE, FastMoE and FasterMoE systems.

TA-MoE: Topology-Aware Large Scale Mixture-of-Expert Training
Chang Chen Min Liz.3, Zhihua Wu', Dianhai Yu', Chao Yang2.55
Center for Data Science, Peking University
School of Mathematics Sciences, Peking University
National Engineering Laboratory for Big Data Analysis and Applications, Peking University *Baidu Inc.
Institute for Computing and Digital Economy, Peking University
charlie chen, chao latao_yan@ppuu edu. cn
limin cww163. com
wuzhihua02, yudi yuzhihuaa2,.udiinaaaaaee www.ddaiiaaaaaiuuuu..oommmmmmm @baiduaaac..mmmo com
Abstract
Sparsely gated Mixture-of-Exper (MoE) has demonstrated its effectiveness in scaling up deep neural networks to an extreme scale. Despite that numerous efforts have been made improve the performance of MoE from the model design system optimization perspective, existing MoE dispatch patterns are still not able to fully exploit the underlying heterogeneous network environments, In this paper, we propose TA-MoE, a topology-awer aware routing strategy for large-scale MoE trainging, from model-system co-design perspective, which dynamically adjust the MoE dispatch pattern according to the network topology. Based on communication modeling, abstract the dispatch problem into an optimization objective and obtain the approximate dispatch pattern under different topologies. On top of that, design topology-aware auxiliary loss, which adaptively route the data to fit in the underlying topology without sacrificing the model accuracy. Experiments show that TA-MoF can substantially outperform its counterparts on various hardware and model configurations, with roughly 1.01x-1.61x, 1.01x- 4.77x, 1.25x-1. 1.54x improvements over the popular DeepSpeed-Mott FastMoE and FasterMoB systems.
Introduction
The scale of model parameters in neural networks has increased from millions trillions in recent years, which promotes model accuracy in many domain, such as language processing 3 4 5 and computer vision 227 23). However, the limited hardware resources, c.g., memory capability and communication bandwidth have constrained the model size to further scale up. To relieve this tension and improve the model performance, Mixture of Expert (MoE) with sparsely gated structure was recently reintroduced [66] 26 25). The core structure of MoE is a group of small "expert" networks and gate network. Guided by the gate result, input data dynamically routed only a sub- sub-group of experts for computation. Compared with dense training methods, the sparsely activated feature of MoE can significantly reduce the computation burden, extend the model size, and achieve higher accuracy 6
13
'Equal Contribution.
Work done during internship Baidu Inc..
Conresponning author.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Since MoE plays vital role in large-scale model training, efficient MoE parallel training has recently received much attention. As one of the most popular MoE training upproaches (Figure . expert parallelism [ 7 distributes experts to different devices, and each device is responsible for a different batch of training samples. Correspondingly, extra global communication is necessary for data exchanges among devices. Recent works aim to increase expert parallelism performance from two aspects. On the one hand, the dynamic pattern of MoE results in severe computation load-imbalance problems that small number of experts may receive, process, and send the majority of data. Several approaches were proposed to make full use of the available experts, such as adding an auxiliary loss (26). controlling expert capacity [1] 77, and optimizing the assignment scheme for On the other hand, global communication is another main obstacle to efficient MoE training. Most of the recent works reduced the communication cost from system perspective, such as computation and communication overlapping [
a balanced load 12 28 22
customized communication
operation acceleration 21 18 and adaptive routing 17
In addition to the continuing efforts made to improve the performance of MoE, there are still two major challenges. With the development of the complicated distributed network environments, the existing even dispatch method may cause network contention in the slowest links, leading to poor communication performance, especially on heterogeneous networks. Although a few early works [9] have proposed methods to dispatch more data to slow links, these methods may make the expert load imbalanced and could influence the model accuracy. Efficient communication demands more delicate dispatch strategies. How to improve the training efficiency without sacrificing the model accuracy is still worth studying. Besides, most of the existing communication optimizations for MoE 21 [7] studied with a specific hardware environment. How to develop methods that can adapt to a variety of hardware environments is also of great practical value.
To tackle these challenges, we design TA-MoE, topology-aware large scale MoE training method that can adaptively adjust the communication volume to fit the underlying network topology. By abstracting the dispatch problem into an optimization objective based on the communication mod- eling, we obtain the approximate dispatch pattern under different topologies. On top of that, an auxiliary topology loss with pattem-related coefficients is proposed, which can dynamically adjust the communication volume without interfering with the model convergence. TA-MoE can also be casily incorporated into the widely used MoE systems, such as DeepSpeed-Men [1 and FastMoE [.
We conduct experiments on various typical network topologies and model configurations. Results show that TA-MoE can substantially outperform DeepSpeed-MoI and FastMoE with roughly 1.0lx- 1.61x speedup and 1.01x-4.77x speedup on different configurations without sacrificing the model accuracy. Compared with the recently proposed Hir gate of FasterMoE, our method can achieve 1. 1.54x speedup on time to convergence. Besides, a more detailed analysis of communication and data dispatch pattern further demonstrates the effectiveness of the proposed data dispatch strategy. The code of TA-MoE is available at: htttps://githubbccom http://tttuuu.ooohhhheeeeee CheCCChhhnhhaaaae Chang/TA-Moo
2 Related Work
Several frameworks have featured sophisticated designs to support efficient MoE training. GShard []] and DeepSpeed-Me [2] subtly composed several einsum operators into the computation of MoE but introduced redundant zero computation and extra memory consumption. FastMoE [ customized the essential computation kernels to improve resource utilization effectively. To further enhance the performance, most of the systems adopted an auxiliary loss (26 to achieve an even dispatch pattern and enforced the number of data processed by each expert below some uniform capacity. Based on these popular implementations, recent works aim to improve the MoE training performance from mainly two aspects: model structure design and communication optimization.
From the perspective of model design, BASE Layer [2 and the work of expert choice routing 28 assigned an equal number of tokens to each expert by delicate designs of the gate. Instead of learning the weight of the gate, Hash Layers [22] adopted an efficient hash function to guide the dispatch. The hybrid structure of PR-MoE 21 improved the parameter efficiency by fixing one shared expert. BaGual a [5] re-distributed the data chunks evenly, damaging the model accuracy. However, almost all of these high-level algorithms are agnostic of the complicated underlying hardware effect on training performance.

---

As for communication optimization, DeepSpeed-Moo [2]] and HetuMoe [18] implemented a hierar- chical all-to-all communication kernel to improve network utilizzatonn Tutle [77] designed adaptive routing techniques coupled with specific network architecture. Despite of these delicate designs, the improvement space of system- level optimization is significantly constrained by the dispatch patterns of MoF. Recently, FasterMoF. [9] made an initial try to take the dispatch pattern into consideration by setting a compulsory ratio of intra-node to inter-node dispatch chunk sizes but sacrificed some model accuracy. In this paper, propose topology-aware routing strategy that enables an effi- cient communication pattern to fit into the underlying topology without sacrificing the convergence performance.
3 Background
3.1
MoE Model Structure
Global Exchange
Global Exchange
Token
Device
Gate Network
Expert Network
Token 0
Aggregate
Token
Device
Gate Network
Token 1
Expert Network
Token 1
Aggregate
Token 2
Token
Device
Gate Network
Expert
Token 2
Aggregate
Token 3
Device 3
Gate Network
Expert Network 3
Token 3
Aggregate
Figure 1: The popular expert parallelism method of MoE.
A MoE layer consists of gate network G and a set of N expert networks E..... Ex-1. E.a....,,x-1.. For the gate module, the softmax activation function widely used, reflecting each expert's 'ss normalized fitness for dealing with an incoming sample. Usually, only the experts with the top & fit scores are selected to process the sample. The final output y of the MoE layer is the aggregation of computed results.
Expert parallelism has been one of the most popular methods in existing MoE training systems 21 As shown in figure [ the N experts evenly assigned P devices, with each device i holding E' = N//P experts E, ((++1)+E-11 Besides, the input tokens are also evenly partitioned over multiple devices with different small batches of the same size S in traditional data-parallel way. For each process, the shape of the dispatched data is jk * S, dI. where d represents the hidden size. Each expert receives a combined batch of tokens (Global Exchange) and then carries out the computation. During the global communication, the number of samples sent to & from process i is Cie, and the shape of the transferred samples is [Cte, d]. Afterward, the expert sends the calculated result back with similar global exchange pattern.
However, the number of the tokens processed by different experts may be highly imbalanced that a small group of experts may receive the majority of data, like Expert 2 in Figure [ Therefore, a load-balance auxiliary loss term lazz (26 is added to the train loss overall loss function:
N-1 m: EG()) lauz Σ (mie * (Cie/S)
The auxiliary loss can dynamically adjust the value of Gi into target & S/N. To further ensure load balance, uniform data process capacity C is set for each expert in many MoE training systems. DeepSpeed-Mer MoE [21] decomposes the expert capacity C evenly into the local capacity for each process and prunes the exchange size by the local capacity directly: Cse I Cie C/P. FastMoE [8] efficiently uses the capacity with two extra all-to-all communication for exchange sizes:
En- Gie C.
3.2 Network Topology
The network environments are very complicated for distributed training on modern GPU clusters. As shown in Figure there are four kinds of typical network topologies: tree and asymmetrical tree. Homogeneous and ring structures are frequently used topologies for the intra-node environment. For a homogeneous structure, devices are always connected by the network
homogeneous, ring, symmetric
3

---

GPU
GPU
GPU
Layort
Layer?
NySwitch
Laynr
CPU
Layero
Layee
GPU
GPU
We characterize the communication cost using the well-known a-3 cost model, where and | represent the fixed communication latency and the inverse bandwidth (i.c., transferring costs of each word), respectively. For convenience, ai and Bi are used to denote the latency and inverse (a) Homogencous
(b) Ring
(c) Symmetric Tree
(d) Asymmetrical Tree
Figure 2: Typical network topologies on modern GPU clusters. (a) A homogeneous node connected
with NVSwitch. (b) A typical ring topology connected with NVLinks [9]]. (c) A 2-layer symmetric tree topology of [2,2]. (d) A 3-layer asymmetrical tree topology of [[2.2].[2]].
with the same bandwidth, e.g., NVSwitch [20] As for the ring topology, it is usually symmetrical. The bandwidths between adjacent devices may differ due to different numbers connected links. The communication of nonadjacent devices has to hop through intermediate devices and the slowest link may become the bottleneck. Hierarchical tree a common topology abstraction for multi-node distributed environments. Compared with the intra-node environment, inter-node links suffer from limited and volatile bandwidth (4~25GGB/s and potentially degrade the communication performance. For convenience, we denote tree topology sub-list are connected by the same switch. For a symmetric tree structure, we use L; to represent the number of the child nodes of each node in layer i. As for asymmetrical tree structure, it is the most common topology for distributed training, which can be very irregular.
nested list where the elements within the same
3.3 Motivation
The existing load-balanced data distribution Equation i unable fully exploit the complicated distributed environments. To demonstrate it, set up experiment topology cluster, where the devices are named 0, (same node) and 0, (same node), respectively, We dispatch 128MB 4 data with two dispatch patterns: (1) even dispatch and (2) uneven dispatch that a greater proportion of data is exchanged with neighbor device. Table shows the detailed dispatch proportions and the corresponding performance. Compared with even dispatch, uneven dispatch improves the overall communication performance by roughly 30%. This is mainly because the communication stress on inter-node links is relieved by transferring a smaller proportion of data. With the variety of distributed network topologies and their continuous development, the existing static even dispatch pattern not effective enough. There is an urgent need for routing strategy that can dynamically adapt the underlying network environments.
[2, 2] symmetric tree
Table 1: The communication
performance [[0, i]] network topology.
Dispatch Pattern
0 <+
0
<+
0 <>>
0 <+ 1
Even
1/4
174
1/4
1/4
Ratio data
Uneven
1/4
1/2
1/8
1/8
Even
144
758
5609
5618
14019
Time (us)
Uneven
144
1492
2835
2861
10765
4 Topology Aware Routing Strategy
In this section, we first abstract the data dispatch problem into an optimization objective based on the communication model. Through some analysis, we obtain the target dispatch pattern under different topologies, which can eliminate the communication bottleneck during MoE training. Guided by the target pattern, we design topology-aware routing loss to adaptively adjust the dispatch volume.
4.1 Communication Model
bandwidth between the i-th and j-th GPU. During the training of MoE, the amount of data transferred
Herre 128MB is used as demonstration, which is the upper-bound bound transfer size of most of the typical MoE
training tasks.

---

from GPU i to & in GPU is Cie * * b, where d b is the transferred element size. To reduce the overheads of multiple send-receives between two GPUs, we merge the multiple small data chunks into larger data chunk for delivery. The total amount of data delivered from GPU i to GPU
Ze=Baj ΣEE Cie d A global data exchange consists of P P peer-to-peer data deliveries, among which the slowest delivery, lower-bound, constrains the final communication performance. Most of the global exchange implementations (2T [8] (4 are designed to approach the lower-bound. Therefore, our ultimate objective function is to minimize the slowest send-receive communication cost:
E-(j+1)-1
min max + Big 453
Σ
Cie d dd+b)
E+
For efficient MoE training, two constraint conditions should be satisfied. First, for any process i, the sent data size, i.c., k * S, should be equal to the sum of received data size of all experts:
k+S
Σ
Ciay Vi E
(O.....,NNI)
Second, to make full use of all the experts and pursue better model accuracy, the data chunks dispatched to each expert should be balanced:
k * S
Σ
Cie Ve f0, N 1}.
(4)
E
P-10
4.2 Model Optimization
To get the target dispatch pattern, we need solve the optimization problem in Equation 2 Nev- ertheless, Equation 2 contains plentiful parameters of specific network, which complicates the solving process. Meanwhile, some irregular topologies, some devices may suffer from quite limited bandwidth when communicating with other devices. According to Equation 2 the experts assigned to these devices may receive quite small dispatch chunk size from the other processes, which may make the experts lack of sufficient data exchanges and lead to expert isolation phenomenon. To tackle these problems, we simplify the optimization problem to accelerate the solving process and smooth the values of oij,, 3:j for approximate result to prevent expert isolation. Since each send-receive communication shares the same a, 3 in homogencous network, the target dispatch chunk size che equal the load-balanced chunk size kes9 In the following part, we focus on the analysis of the optimization problem under heterogeneous topologies.
On n--laye symmetric tree topology, for any device i, all the devices can be split into sub-groups of G = {G < n). G; is the group of devices whose shortest path from device are across t switches. Multiple hops in cross-switch communication will suffer from extra overheads and the most limited bandwidth in the hops dominates the final bandwidth. Therefore, we can simplify Eizs I(EGG))a
Eiee ((
the original Big into value:
aT (L):((=11)/>= 3 (T Lu)=(LI-1)/2² which can
precisely characterize the underlying topology and eliminate the noise of profiling. Then we get hierarchical matrix á B:
d:j ΣnG E G aa. B.g = Σ T( E GUE
(5)
Take equation and optimization constraint conditions, we simplify the optimization problem as follows:
E+(j+1)
min comm
min max(a + B.j 1.j
Σ
Cie d +b))
c= E=Enj
S
s.t.
Σ
Cie * d b, Ve (o...., 1},
(6)
:E(O....P-1)
k 5 =
Σ
Cie * d bi Vi (0.... P-1}.
of Ne(0....N11I
c c20
5

---

Then, we can get the near optimal solution after omitting the small latency term:
k S
Eie
E Ej + 1c(
The optimal data distribution of the above min-max problem is only related to the bandwidth: the volume of Eie is linear to the bandwidth. The rationale behind the optimal result is that higher bandwidth links should bear more loads for an overall communication balanced workflow. The ring topology also shows a hierarchical characteristic and the solution for ring topology has the same pattern symmetric trees.
Under some irregular asymmetric topology, the optimal result of Equation 2 may result in some experts assigned to the devices of most limited bandwidth in lack of data exchanges with the rest of the devices. Compared the data distribution of other experts, data chunks from local devices occupy a larger proportion of received data chunks in those isolated experts. For the fairness among the experts, we transform the asymmetric topology into a symmetric one by merging the separate nodes into the close symmetric sub-trees. For example, [[2,2][2]] in figure 2(d) can be merged as symmetric structure [[2,2.2]]. After that, we can optimize the lower bound of communication as the symmetric structure.
4.3 Routing Strategy
Once getting the target data dispatch volumes among processes, use it guide the MoE training. To not sacrifice the model accuracy, instead of setting a compulsory dispatch ratio directly, design topology-awe awarre adaptive routing loss.
pt Norm(1/2), Hopo N P
Σ (pie mie (Gt)/S) Vi E t0, P-13 13
(8)
As shown in Equation
g we set a penalty weight p; as the adjustment coefficients for the topology loss liopo of each process 2. We set the normalized 1/2, as p; to make sure Cse approximates the value of cie- Normalization functions that enlarge the penalty of the low- www-bandwidth transfer, e.g., softmax, are also preferable. As shown in the calculation of 12 Ropo the data dispatched to & with limited bandwidth will suffer from larger penalty weight Pie- Despite of these penalty modifications, auxiliary loss occupies a small proportion of the final loss, and the value of the auxiliary loss decreases when experts' number scales up. Therefore, our final topology loss is expanded N P times to keep the magnitude of loss value.
There are a number of advantages of the proposed topology-aware strategy. On the one hand, compared setting a compulsory dispatch ratio, the proposed loss can adjust the communication volume to fit in the underlying topologies in mild way without damaging the convergence. A compulsory dispatch ratio has high potential to overwhelm the influence of the train loss and sacrifice the model accuracy. With the topology loss, the train loss can still dominate in the final loss value. As a result, the dispatch results are mainly influenced by the train loss for better model accuracy. Besides, the topology-awan awarr strategy has more potential to utilize the token information for efficient sparse training. Guided by the topology-awan loss, the tokens nearby are more likely to be processed by the same expert. Since the correlation of adjacent tokens contains the vital information in sparse attention [10]. the experts in sparsely gated routing structure may also more likely to extract important information from adjacent tokens.
In addition, the topology-awart aware loss can be easily incorporated existing MoE systems, such as DeepSpeed-Mon and FastMoE. Taking FastMoE as an example, one just needs to directly replace the popular
loss laus with the proposed ropology-aware loss ltopo- Since local capacity threshold Cie is adopted in DeepSpeed-Mer sizes to be consistent with the proposed dispatch pattern by setting Cte to be proportional to the target data chunk sizes C.. Instead of padding the data chunks with extra zeros to be the same size as DeepSpeed- one all-to-all communication added to get the information of send-receive data chunk sizes and dispatch data chunks according to the sizes.
for load balance, one can modify the local capacity

---

Evaluation
Experiment setup To demonstrate the effectiveness TA-M we carry out series of experi- ments on three typical NVIDIA GPU clusters with different network topology. and some represen- tative model configurations. Table 2 lists the cluster settings. The testbed is three typical clusters from the PaddleCloud platform. For cluster A, each node consists of 8 NVIDIA Tesla 40GB A A00 GPUs connected with NVSwitch, which shows high performance for both computation and network communication. Clusters B and C are equipped with 8 NVIDIA Tesla 32GB V100 GPUs in each node. The nodes in cluster are connected by the same switch, while cluster is composed of a large number of servers and switches that are interconnected through an internal network infrastructure. Besides, the software configurations are set as CUDA 1 11.0, NCCL 2.8.4 and CUDA 11.1, NCCL 2.8.3, for cluster A and cluster B, C, respectively.
Table 2: Cluster Setting.
Clusters
GPU
Intra-Node Inter-Node
Symmetric Same switch
A
40G-A100
NVSwitch
100GB/s RoCE/4
X
B
1200
NVlink 100GB/s RoCE/S
V
C
32G-V100
NVlink 100GB/s RoCE/8
X
Since the transformer architecture is the most common base structure of MoE, we focus the experi- ments problems related to wansformer-basee model, with GPT-3 Medium B] as the base model and multi-layer perception as the expert. In our experiments, the number of the experts are chosen among {8, 16, 32, 48, 64) with each device deployed with one expert. Both the Switch top-1 (7) and the GShard top-2 gates [11] are tested with the weight of auxiliary loss set as 1.0. For the consistency of the experiment, we implement the models by single framework Paddle [2] and train on the open-source openwebtext2 dataset [[] More detailed specifications of model settings can be found in Table 3 To be more general, we also add the tests of MoE on Swin Transformer based MoE tasks in Section A.3 of the Appendix. These results further demonstrate the effectiveness of the proposed topology-awe aware routing algorithm on different model architectures.
Table 3: Detailed specifications of the GPT models.
Gate
Layers
Hidden size Intermediate size Batch size
Data type
Capacity factor Clusters
Switch
12
I024
4096
6
FPI6
1.0
A
GShard
12
1024
2048
6
FP16
2.0
A
Switch & GShard
12
1024
2048
4
FP32
1.2
B
Switch & GShard
12
1024
2048
4
FP32
1.2
C
Methodology We incorporate TA-MoE into widely used DeepSpeed-Mol 21 and FastMoF g implementations. Because TA-MoE modifies the gate structure, we first compare the validation loss w.r.t. steps to ensure that TA-MoE will not interfere with the convergence on various model scales. On top of that, we test the overall throughput and the speedup of TA-MoE over these two classical baselines. To be more comprehensive, we also compare with the recently proposed FasterMoB Hir gate [9] on the metric of time to convergence performance. Besides, a detailed analysis of communication costs, as well the distribution of the dispatch are also given.
Accuracy and Performance Like some well- recognized related works, e.g., BASE layer and DeepSpeed-MoE, we have depicted the validation performance of every fixed interval step Figure 3 as the comparison metric. We first compare the validation loss w.r.t steps of TA-MoE and the repre- sentative FastMoE cluster C. As shown in Figure 3. the loss curves of TA-MoE and FastMoE are con- sistent converge under different training scales of 8 experts to 48 experts. These results demonstrate that
Daseline EB TA-MoE E8
Daseline E32
TA.MoE E32 TA MoE, E48
Volidation Loss
MAE E11
3.0
0
20000
40000
60000
80000
100000
A Cloud Platform of Baidu Inc.
Figure 3: Validation loss w.r.t. steps.
7

---

the TA-MoE can adaptively adjust the data dispatch volume without sacrificing the model accuracy. To be more comprehensive, we also add the perplexity (PPL) metric Table of the Appendix, which gives similar effective results.
Switch Gate
GShard Gate
1e5
Throughpu (taker 9
Baseline
1,05x
1.15x
PostMoE+TA
1.11x
DeepSpend Baseline
1.13x
Cluster A
Deep Speed TA-MoE
1.13x
4-
2
3-
1.61 47H
2. 2.40x
1.26x
1.54x
00M
OOM
1.13x
1.19k
Througheer Bruughattti (tolonnis)]
1.2
1.23x
1.23x
1.0-
1.25x
1.49x
Cluster B
1.0
1.32x
1.60x
0.8
0.8
0.6
121
1.16x
1.30x
1.58x
1.27%
1001
1.17x
0.4
1.02%
1.20x
0.2
E8
E16
E32
E48
E64
E16
E32
E48
E64
1.13x
1.18x
1.2
Throughus (tokangjs)
1.22x. 1.22x
1.25x
1.0
Cluster C
1.29x
0.8
1.31x
1.30x
1.24x
136x
1.38x
1.71x
1.01
0.4
4.77x
E16
E32
E48
E64
E16
E32
E48
E64
Figure 4: Performance of TA-MoE over DeepSpeed-Mon and FastMoE under different hardware and model settings.
To demonstrate the advantages of TA-MoE in terms of training performance, compare it with both DeepSpeed-Mof and FastMoE on various hardware and model configurations. The performance indica- tors including throughout (tokens/s) and the speedups are depicted in Figure 4 is clear that TA-MoE can bring significant performance improvements its counterparts under almost all the configu- rations. When compared with the DeepSpee-MoE, the achieved speedup is about 1. 1.61x-1.61x. As for FastMoF, the performance improvement is around 1.01x-4.77x. More speedups are achieved on Fast- MoE than DeepSpeed- MoF. This is because TA-MoR. has more dynamic dispatch pattern and larger ad- justable space based on FastMoE.
RosterMoE E32
TA-MFF_F3 F32
validation Loss
3.5
1.25x
3.1
1.47x
2.9
1.54x
2.7
500
1000
1500
2000
2500
3000
Breemin)
Figure 5: Comparison with FasterMoE.

---

It is also observed that more improvements are obtained on Cluster C. which can reach 4.77x for some cases, due to the relief of its serious network contention of cross-switch communication. In addition, the comparison of the results of Switch and GShard gate reveals that TA-MoR can behave better in adjusting larger volume of data.
To be more comprehensive, we make further comparisons with the recently proposed FosterMoE [2]. Because the compulsory dispatch strategy of FasterMoE affects the convergence, take the validation loss w.r.t time as the comparison metric. Clusters C are selected as the representative testing clusters. As shown in Figure TA-MoE can converge faster than FasterMoE. We evaluate the time to reach the validation loss values of 3. 2.9, and 2.8, and TA-MoE can converge faster by about 1.25x, 1.47x and 1.54x. The results further verify that the proposed adaptive routing loss is more effective than the compulsive dispatch method, which sacrifices model accuracy for training performance.
To better show the effects of the dynamical dispatch strategy, we further analyze communication and computation cost, and the distribution of data dispatch on cluster As shown in Figure 6(a)] thanks to the proposed TA-MoE strategy, the communication cost is reduced rapidly, with roughly 1.16x 6.4x speedups. It also observed that the maximum speedup is achieved for 32 experts on four cross- nodes. This is because the four nodes are deployed under four different switches, and cross-switch links severely bottleneck the data exchange. Once the tension is relieved, the obtained benefits can be dramatic. In addition, we visualize the dispatch pattern of an example with 64 expert by depicting the number of the tokens of Rank 0-7 sending to other ranks. More details of the dispatch distributions for other expert scales are attached to Section A.2 the Appendix. In Figure [6bb] as expected, most of the data of Rank 0-7 are dispatched to low-overheads nearby ranks, which further verifies the effectiveness of adaptive topology- aware loss.
Communication Analysis
Computation and Communication Cost (ms)
1e3
>
6.40x
Baseline
6
FastMoE Comp Comm
Dispatched Taken Number
1.0
To-oop
0.8
Rank- 2
Parkk
Rank
BankeS
3-
1.51x
1.16x
0.0
E8
E48
10
30
50
(a) Breakdown of communication and computation.
(b) Distribution of data dispatch of Rank 0-7.
Figure 6: Analysis of communication and computation cost and the distribution of data dispatch.
Conclusion
In this paper, topology-aware routing strategy, TA-MoE, was proposed to stress the mismatch between the data dispatch pattern and the network topology. Based on communication modeling, we abstract the dispatch problem into optimization objective and obtain the approximate dispatch pattern under different topologies. On top of that, a topology-awe aware auxiliary loss was designed, which can adaptively route the data to fit in the underlying topology without sacrificing the model accuracy. Experiments show that the proposed method can substantially outperform its counterparts on variety of the hardware and model configurations, with roughly 1.00x-1.61xx 1.01x-4.77x, 1. 1.25x-1. 1.54x improvements over the popular DeepSpeed-MoF FastMoE and FasterMoE. systems. In the future, we plan take more delicate communication operators into consideration of dynamic dispatch pattern and extend our work to more hardware environments.
7 Acknowledgement
This work was supported in part by National Key R&D Program of China 0021ZD0110500

---

References
[1] Openwebtext2. http://:/ggggttttttto.eoaaoooaaa 2022.
[2] Baidu. PaddiePaddle. attps://ggtabbb.oo///aaaapaapppp
2022.
[3] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Vone Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 1877-1901. Curran Associates, Inc., 2020.
[4] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.022: 02220...2111, 2022.
[5] Jacob Devlin, Ming- wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio, editors, Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume (Long and Short Papers), pages 4171-4186. Association for Computational Linguistics, 2019.
[6] Nan Du, Yanping Huang, Andrew M. Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, Barret Zoph, Liam Fedus, Maarten Bosma, Zongwei Zhou, Tao Wang. Yu Emma Wang, Kellie Webster, Marie Pellat, Kevin Robinson, Kathleen S. Meier-Hellstern, Toju Duke, Lucas Dixon, Kun Zhang, Quoc V. Le, Yonghui Wu, Z. Chen, and Claire Cui. GLaM: Efficient scaling of language models with mixture-of-experts. ArXiv, 12.02112.0695... 2021.
[7] William Fedus, Barret Zoph, and Noam Shazeer. Switch Transformers: Scaling to trillion parameter models with simple and efficient sparsity. CoRR, hts/2101.03961.. 2021.
[8] Jiaao He, Jiezhong Qiu, Aohan Zeng, Zhilin Yang. Jidong Zhai, and Jie Tang. FastMoB: A fast mixture-of-expert training system. CoRR, abs/2 13262, 2021.
[9] Jiaao He, Jidong Zhai, Tiago Antunes, Haojie Wang, Fuwen Luo, Shangfeng Shi, and Qin Li. FasterMoE: modeling and optimizing training of large-scale dynamic pre-trained models. In Proceedings of the 27th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming, pages 120-134, 2022.
[10] Torsten Hoefler, Dan Alistarh, Tal Ben-Nun, Nikoli Dryden, and Alexandra Peste. Sparsity in deep learning: Pruning and growth for efficient inference and training in neural networks. Journal of Machine Learning Research, 22(241):1-124, 2021.
[11] Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. GShard: Scaling giant models with condi- tional computation and automatic sharding. In International Conference on Learning Represen- tations, 2021.
[12) Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. BASE layers: Simplifying training of large, sparse models. CoRR, abs/2103.16716. 2021.
[13] Junyang Lin, An Yang, Jinze Bai, Chang Zhou, Le Jiang, Xianyan Jia, Ang Wang, Jie Zhang, Yong Li, Wei Lin, Jingren Zhou, and Hongxia Yang. M6-10T; A sharing-delinking paradigm for efficient multi-trillion parameter pretraining. CoRR, abs/2110.038888, 2021.
[14] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. CORR, abs/2103. 14030, 2021.
10

---

[15] Zixuan Ma, Jiaao He, Jiezhong Qiu, Huanqi Cao, Yuanwei Wang. Zhenbo Sun, Liyan Zheng, Haojie Wang, Shizhi Tang, Tianyu Zheng, al. BaGuaLu: targeting brain scale pretrained models with over 37 million cores. In Proceedings of the 27th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming, pages 192-204, 2022.
[16] Saeed Masoudnia and Reza Ebrahimpour. Mixture of experts: literature survey. Artificial Intelligence Review, -2:275-293, 2012.
[17] Microsoft. Tutel: An efficient mixture-of-experts implementation for large dnn model training. httpp///ggtaabbbccmom///crrrrrts 2021.
[18] Xiaonan Nie, Pinxue Zhao, Xupeng Miao, and Bin Cui. HetuMoE: An efficient trillion-seae scale
mixture-of-expert distributed training system. arXiv preprint arXiv:2203. 2022.
[19] NVIDIA. NVLINK. attps://xww.nvidtt http://twww////ooooo///.././///a
2022.
[20] NVIDIA. NVSWITCH: The world's highest- on-nod switch. btttps:/iimage 2022.
nvidia. cww
[21] Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi, Ammar Ahmad Awan, Jeff Rasley, and Yuxiong He. DeepSpeed-MoE: Advancing mixture- of-experts inference and training to power next-generation scale. CoRR, abs/2201.05596, 2022.
[22] Stephen Roller, Sainbayar Sukhbaatar, arthur szlam, and Jason Weston. Hash layers for large sparse models. In M. Ranzato, Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 17555-17566. Curran Associates, Inc., 2021.
[23] Carlos Riquelme Ruiz, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. In Beygelzimer, Y. Dauphin, P. Liang, and Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[24] Aashaka Shah, Vijay Chidambaram, Meghan Cowan, Saeed Maleki, Madan Musuvathi, Todd Mytkowiez, Jacob Nelson, Olli Saaríkivi, and Rachee Singh. Synthesizing collective munication algorithms for heterogeneous networks with TACCL. CoRR, abs/2111.00 11.04667, 2021.
[25] Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, Penporn Koanan- takool, Peter Hawkins, HyoukJoong Lee, Mingsheng Hong, Cliff Young, Ryan Sepassi, and Blake Hechtman. Mesh-TensorFlowe Deep learning for supercomputers. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, NIPS' 18, page 10435-10444, Curran Associates Inc., 2018.
[26] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc V. Le, Geoffrey E. Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of- experts layer. CoRR, 04x/1701.065344 2017.
[27] Bichen Wu, Chenfeng Xu, Xiaoliang Dai, Alvin Wan, Peizhao Zhang, Masayoshi Tomizuka. Kurt Keutzer, and Peter Vajda. Visual transformers: Token-based image representation and processing for computer vision. CoRR, abs/2006.03677, 2020..
[28] Yan-Quan Zhou, Tao Lei, Han-Chu Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M. Dai, Zhifeng Chen, Quoc V. Le, and James Laudon. Mixture-ofeexperts with expert choice routing, ArXiv, 0002220..09388, 2022.

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] Please see the conclusion section
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you including theoretical results...
(a) Did you state the full set assumptions all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed reproduce the main experimental results (either in the supplemental material or as URL)? [Yes]
(b) Did you specify all the training details (e.g.. data splits, hyperparameters, how they were chosen)? [Yes] Please the experiment setup part in section 5.
(c) Did you report error bars (e.g.. with respect to the random seed after running experiments multiple times)? [Yes]
(d) Did you include the total amount compute and the type resources used (e.g.. type of GPUs, internal cluster, or cloud provider)? [Yes] Please see the experiment setup part section
4. If you are using existing assets (c.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [Yes]
(c) Did you include any new assets either in the supplemental material or as a URL?
[No]
(d) Did you discuss whether and how consent obtained from people whose data you'l 're using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identi- fiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent participant compensation? [N/A]
12

---

A Appendix
A.1 Perplexity Evaluation Results
Table 4: Perplexity Evaluation Result. To further validate the model convergence performance, we list the perplexity (PPL) at 10w step (near days) of GPT-Medium (12 layers, hidden size 1024, intermediate size 2048, GShard, Capacity factor 1.2) with different expert numbers on the openwebtext2 dataset. Combined with the information in Figure of the paper, we can find that TA-MoE does not influence the convergence of the model when compared with the well-known FastMoE baseline.
Expert Scale
TA-MoE Valid PPL Baseline Valid PPI.
17.97
18.12
16
15.18
15,39
32
13.37
13.53
48
12.55
12.49
A.2 Data Dispatch Distribution
Figure 7 further elaborates the data dispatch patterns of GPU rank 0-7 (within node) on 8, 16, 32, 48 experts. We also list the Rank-O dispatch information of FastMoF with even distribution method the baseline. On single node topology, each rank prefers to dispatch data to the expert within node. The topology influence on the data dispatch preference is relatively small, because the bandwidth variance within a node is small. Multi-node topology results show a consistent "iade"r-like" distribution trend that the ranks within node has high preference to dispatch the data to intra-node rank group, instead of transferring data to inter-node ranks. These results verify the effectiveness of the proposed adaptive topology-aware method.
1e3
Baseline
Dispatched Taken Numbee
2.0
Rank-0
Number
Rank
Rank-1
Pank-2
1.5
Rank-3
Rank-3
Rank-4
0.8
Rank-5
1.0
Mank-7
Ranke6-
0.6
Ranki6
Rank-7
0.4
0.5
0.2
0.0
0.0
Rank of Exports
Rani Experts
1c3
1c3
1.2
Baseline
Baseline
Dispatched Token Number
1.0-
Rank-1
Dispatched Token Number
0.8
Pank-1
Rank-2
Rank:2
Rank-4
0.6
RanksS
Rank-5
0.6
Rank-7
Pank-?
0.4-
0.2
0.0
0.0
10
15
20
25
10
20
40
Biperts
Figure 7: The data dispatch distribution of rank 0 to 7 on 16, 32, 48 GPUs and experts.
13

---

1.20x
FastMoE
2.0
TA-MoE
Time cost/iter (s)
1.18x
1.5
1.0
0.5
0.0
E16
E32
Figure 8: Speedup of TA-MoE over FastMoE on Swin Transformer Based Model.
A.3 Tests on Swin Tranformer Based MoE Tasks
To further validate the generality of TA-MoE, we also carry out an experiment of Vision Tasks on ImageNet-Ik dataset. The well-known vision transformer architecture Swin Transformer [4] is picked as the base model. The detailed model configurations are listed in Table 5 We evaluate the speedup on Cluster with 16 and 32 GPUs an illustration, where GPUs configurations represents symmetric tree topology and 32 GPUs configurations represents asymmetric tree topology. As shown in Figure [ we can achieve 1.18x and 1.20x speedup when compared with FastMoE on 16, and 32 GPUs, respectively.
Table 5: Detailed specifications of the Swin-Transformen model.
Capacity factor
Name
Layers
Gate
Stage
Stage
Stage 3
Stage
4x4,
concat 2x2,
concat 2x2,
concat 2x2,
Swin Transformer
96-d, LN
192-d, LN
384-d,
768-d, LN
12
GShard
[win.sz. 7x7,
[win.sz. 7x7,
(win.sz. 7x7,
[win.sz. 7x7,
1.2
vl
dim 96,
dim 192,
dim 384,
dim 768,
head3] x2
head6) x2
head12] x6
head324 x2
14


Reviews:
[{"review_id": "xGdY48er99r", "reviewer": "Reviewer_QoUG", "summary": "Sparsely gated Mixture-of-Expert (MoE) plays a vital role in large-scale model training but suffers from both load imbalance and global communication. In addition, the existing even dispatch approach may cause network contention and worsen the previous challenges. This work proposed a topology-aware large-scale MoE training method, called TA-MoE, that can adapt communication volume to fit the underlying network topology without interfering with the model convergence. The key ideas are abstracting the dispatch problem as a communication cost optimization problem and then adding an auxiliary loss with pattern-related coefficients. Experiments show that TA-MoE provides up to 1.61x speedup and 4.77x speedup over DeepSpeed-MoE and FastMoE without accuracy loss.", "questions": "Please refer to the above weakness part.", "limitations": "The authors have adequately addressed the limitations and potential negative societal impact of their work.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n+ this work tried to tackle a very significant and interesting challenge in MoE system: network topology may worsen the communication and load balance problems during the dispatch in MoE.\n+ the paper is well organized and easy to follow\n+ the proposed TA-MoE method is simple and effective: extensive experiments show that TA-MoE is able to offer noticeable speedup over the state-of-the-art under different hardware  and model configurations.\n\nWeaknesses:\n- the experiments are mostly doen with GPT models; it would be better to have models with different neural architectures in the evaluation benchmark. It is unclear how TA-MoE works on other MoE using models other than GPTs.", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "OZBQk6it-V7", "reviewer": "Reviewer_VbEU", "summary": "The paper proposes a new algorithm to improve training efficiency of Mixture of Experts models in a distributed training setting by exploiting the network topology information. To achieve this, the authors propose a new auxiliary loss term incorporating communication bandwidth to encourage tokens to be routed to closer nodes rather than further nodes. By applying this new algorithm, authors claim that they could achiever faster throughput (1.01x - 4.77x) without losing accuracy on their several different clusters. As a result, they show a faster wall-clock time convergence.", "questions": "There are a few questions that need to be answered as well. First, it is not clear that how capacity and overflow tokens are handled in the proposed algorithm. They are known to be critical factors for the successful MoE model trainings, but not much details about those are included in the paper.\nSecond, it is not clearly shown how the locality preferred routing impacts the training compared to the global routing. Maybe, it could be useful to have a few data points such as forced local routing. It is unclear how experts can have expertise while trained by local tokens more. It might be useful to see how different topology affects expertise of different experts.\nLastly, a few details are missing (number of layers, datasets, model architecture). Figure 3 and Figure 6 (b) are hard to read.", "limitations": "This paper is focusing on the computation algorithm itself. So, it might not have direct societal impact.", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The communication overhead is one of the major issues for the MoE model training and this paper proposes a new method to deal with this problem naturally. Given the increased usage of MoE model technology, this is a timely work. Having a soft guidance seems like a good idea not to hurt the original training dynamics while encouraging locality of token routing. And, as authors mentioned, there have not been this kind of topology aware loss terms before as far as I know.\nHowever, there are a few missing details about model configurations and algorithms asked in the question section. And, the overall speed gain is minor.", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Ebhq9jKib8R", "reviewer": "Reviewer_EndF", "summary": "This paper addresses the problem of MoE routing under the cases of different network topologies by allocating another abstraction layer for the topology and designing an auxiliary objective to optimize. Experiments show very good improvement in terms of speed compared to strong baselines.", "questions": "As said, while many (if not most, in my humble opinion) AI people are not working near the system level to address the same problems as in this paper, e.g., they simply use AWS, Azure, Colab, or other cloud infrastructure for fast prototyping or deployment, I wonder how much improvement this solution can make with different cloud platforms? In other words, have “cloud” people already seen those heterogeneous infrastructure problems and addressed that at the cloud provision/hypervisors level? The paper presents the case of PaddleCloud but how about others? ", "limitations": "Maybe not very relevant since the paper addresses the system-related level and thus is hard to judge those impacts. ", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength: \n\n1. The paper offers an important contribution to the AI community at the system level, which is probably not difficult to approach for many people working in this field. In fact, in my humble opinion, not so many AI people have the opportunity to access detailed hardware information as cloud users such as with Azure or AWS. \n\n2. The experiments show very good improvement over strong baselines. System analysis is clearly presented. \n\nWeakness\n\n1. The paper addresses the system level. However, since it claims a significant boost of speed without sacrificing the model accuracy, it needs to show the accuracy, e.g. at least the LM-related one with NLP-related metrics. \n\n2. Line 240, which claims \"without loss of generality\", is probably too strong. My suggestion is if the solution is good, with the current hardware settings, the authors can run current codes for other many applications of which codes are available to further solidify their claims. \n\n3. Likewise, why not show the data dispatch distribution of other ranks but only rank 0? If space is limited, appendix space is always there. \n\n4. In the era of GPUs and large data, the motivation is led by demonstrating only 128MB of data is probably inefficient. Probably at least some GBs, or even stronger in a combination with different types of data would make a stronger motivation. \n\n5. No code is provided. \n", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}]

Author rebuttal:

Decision: Accept