TA-MoE: Topology-Aware Large Scale Mixture-of-Expert Training
Anonymous Author(s)
AffiliationAddress email
Abstract
Sparsely gated Mixture-of-Expert (MoE) has demonstrated its effectiveness inscaling up deep neural networks to an extreme scale.
Despite that numerous effortshave been made to improve the performance of MoE from the model design orsystem optimization perspective, existing MoE dispatch patterns are still not able tofully exploit the underlying heterogeneous network environments.
In this paper, wepropose 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 communicationmodeling, we abstract the dispatch problem into an optimization objective andobtain the approximate dispatch pattern under different topologies.
On top ofthat, we design a topology-aware auxiliary loss, which can adaptively route thedata to fit in the underlying topology without sacrificing the model accuracy.
Experiments show that TA-MoE can substantially outperform its counterpartson various hardware and model configurations, with roughly 1.01x-1.61x, 1.01x4.77x, 1.25x-1.54x improvements over the popular DeepSpeed-MoE, FastMoE and
FasterMoE systems.
Introduction
The scale of model parameters in neural networks has increased from millions to trillions in recentyears, which promotes model accuracy in many domain, such as language processing [3, 4, 5] andcomputer vision [26, 22].
However, the limited hardware resources, e g., memory capability andcommunication bandwidth, have constrained the model size to further scale up.
To relieve this tensionand improve the model performance, Mixture of Expert (MoE) with a sparsely gated structure wasrecently reintroduced [15, 25, 24].
The core structure of MoE is a group of small "expert" networksand a gate network.
Guided by the gate result, input data is dynamically routed to only a sub-groupof experts for computation.
Compared with dense training methods, the sparsely activated featureof MoE can significantly reduce the computation burden, extend the model size and achieve higheraccuracy [6, 7, 11, 13].
Since MoE plays a vital role in large-scale model training, efficient MoE parallel training hasrecently received much attention.
As one of the most popular MoE training approaches (Figure 1),expert parallelism [11, 7] distributes experts to different devices, and each device is responsible fora different batch of training samples.
Correspondingly, extra global communication is necessaryfor 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 a small number of experts may receive, process, and send the majorityof data.
Several approaches were proposed to make full use of the available experts, such as addingan auxiliary loss [25], controlling expert capacity [11, 7], and optimizing the assignment scheme fora balanced load [12, 27, 21].
On the other hand, global communication is another main obstacle toefficient MoE training.
Most of the recent works reduced the communication cost from a systemperspective, such as computation and communication overlapping [9], customized communicationoperation acceleration [20, 17], and adaptive routing [16].
In addition to the continuing efforts made to improve the performance of MoE, there are still twomajor challenges.
With the development of the complicated distributed network environments, theexisting even dispatch method may cause network contention in the slowest links, leading to poorcommunication 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 loadimbalanced and could influence the model accuracy.
Efficient communication demands more delicatedispatch strategies.
How to improve the training efficiency without sacrificing the model accuracy isstill worth studying.
Besides, most of the existing communication optimizations for MoE [20, 16] arestudied with a specific hardware environment.
How to develop methods that can adapt to a variety ofhardware environments is also of great practical value.
To tackle these challenges, we design TA-MoE, a topology-aware large scale MoE training methodthat can adaptively adjust the communication volume to fit the underlying network topology.
Byabstracting the dispatch problem into an optimization objective based on the communication modeling, we obtain the approximate dispatch pattern under different topologies.
On top of that, anauxiliary topology loss with pattern-related coefficients is proposed, which can dynamically adjustthe communication volume without interfering with the model convergence.
TA-MoE can also beeasily incorporated into the widely used MoE systems, such as DeepSpeed-MoE [20] and FastMoE[8].
We conduct experiments on various typical network topologies and model configurations.
Resultsshow that TA-MoE can substantially outperform DeepSpeed-MoE and FastMoE with roughly 1.01x1.61x speedup and 1.01x-4.77x speedup on different configurations without sacrificing the modelaccuracy.
Compared with the recently proposed Hir gate of FasterMoE, our method can achieve1.25x-1.54x speedup on time to convergence.
Besides, a more detailed analysis of communicationand data dispatch pattern further demonstrates the effectiveness of the proposed data dispatch strategy.
Related Work
Several frameworks have featured sophisticated designs to support efficient MoE training.
GShard[11] and DeepSpeed-MoE [20] subtly composed several einsum operators into the computation of
MoE but introduced redundant zero computation and extra memory consumption.
FastMoE [8]customized the essential computation kernels to improve resource utilization effectively.
To furtherenhance the performance, most of the systems adopted an auxiliary loss [25] to achieve an evendispatch pattern and enforced the number of data processed by each expert below some uniformcapacity.
Based on these popular implementations, recent works aim to improve the MoE trainingperformance from mainly two aspects: model structure design and communication optimization.
From the perspective of model design, BASE Layer [12] and the work of expert choice routing [27]assigned an equal number of tokens to each expert by delicate designs of the gate.
Instead of learningthe weight of the gate, Hash Layers [21] adopted an efficient hash function to guide the dispatch.
The hybrid structure of PR-MoE [20] improved the parameter efficiency by fixing one shared expert.
BaGuaLu [14] re-distributed the data chunks evenly, damaging the model accuracy.
However, almostall of these high-level algorithms are agnostic of the complicated underlying hardware effect ontraining performance.
As for communication optimization, DeepSpeed-MoE [20] and HetuMoe [17] implemented a hierarchical all-to-all communication kernel to improve network utilization.
Tutle [16] designed adaptiverouting techniques coupled with a specific network architecture.
Despite of these delicate designs, the improvement space of system-level optimization is significantly constrained by the dispatch patternsof MoE.
Recently, FasterMoE [9] made an initial try to take the dispatch pattern into considerationby setting a compulsory ratio of intra-node to inter-node dispatch chunk sizes but sacrificed somemodel accuracy.
In this paper, we propose a topology-aware routing strategy that enables an efficient communication pattern to fit into the underlying topology without sacrificing the convergenceperformance.
Background
[Figure]
3.1 MoE Model Structure
[Figure]
A MoE layer consists of a gate network G and a set of N expert networks E 0 , . . .
, E N − 1 .
For the gatemodule, the softmax activation function is widely used, reflecting each expert’s normalized fitness fordealing with an incoming sample.
Usually, only the experts with the top k fit scores are selected toprocess 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 [20, 8].
As shown in figure 1, the N experts are evenly assigned to P devices, with each device i holding
E = N/P experts E i ∗ E , . . .
, E ( i +1) ∗ E − 1 .
Besides, the input tokens are also evenly partitioned overmultiple devices with different small batches of the same size S in a traditional data-parallel way.
Foreach process, the shape of the dispatched data is [ k ∗ S, d ] , where d represents the hidden size.
Eachexpert receives a combined batch of tokens (Global Exchange) and then carries out the computation.
During the global communication, the number of samples sent to E e from process i is c ie , and theshape of the transferred samples is [ c ie , d ] .
Afterward, the expert sends the calculated result backwith a similar global exchange pattern.
However, the number of the tokens processed by different experts may be highly imbalanced thata small group of experts may receive the majority of data, like Expert 2 in Figure 1.
Therefore, aload-balance auxiliary loss term l aux [25] is added to the train loss as an overall loss function:
[Equation]
The auxiliary loss can dynamically adjust the value of c ie into target k ∗ S/N .
To further ensureload balance, a uniform data process capacity C is set for each expert in many MoE trainingsystems.
DeepSpeed-MoE [20] decomposes the expert capacity C evenly into the local capacity for each process and prunes the exchange size by the local capacity directly: c ie ≤ C ie = C/P .
FastMoE [8] efficiently uses the capacity with two extra all-to-all communication for exchange sizes:
3.2 Network Topology
The network environments are very complicated for distributed training on modern GPU clusters.
Asshown in Figure 2, there are four kinds of typical network topologies: homogeneous, ring, symmetrictree and asymmetrical tree.
Homogeneous and ring structures are frequently used topologies for theintra-node environment.
For a homogeneous structure, devices are always connected by the network
[Figure]
GPU
[Figure]
GPU GPU GPU (a) Homogeneous
GPU
GPU GPU GPU
GPU
GPU
GPU GPU (b) Ring
GPU
GPU GPU GPU
CPU
CPU Layer
Layer (c) Symmetric Tree
GPU
GPU GPU GPU GPU
GPU
CPU CPU
CPU Layer
Layer 1 Layer
[Figure]
with the same bandwidth, e.g., NVSwitch [19].
As for the ring topology, it is usually symmetrical.
The bandwidths between adjacent devices may differ due to different numbers of connected links.
The communication of nonadjacent devices has to hop through intermediate devices and the slowestlink may become the bottleneck.
Hierarchical tree is a common topology abstraction for multi-nodedistributed environments.
Compared with the intra-node environment, inter-node links suffer fromlimited and volatile bandwidth (4 ∼ 25GB/s) and potentially degrade the communication performance.
For convenience, we denote a tree topology as a nested list where the elements within the samesub-list are connected by the same switch.
For a symmetric tree structure, we use L i to represent thenumber of the child nodes of each node in layer i .
As for an asymmetrical tree structure, it is themost common topology for distributed training, which can be very irregular.
3.3 Motivation
The existing load-balanced data distribution of Equation 1 is unable fully exploit the complicateddistributed environments.
To demonstrate it, we set up an experiment on a [2, 2] symmetric treetopology cluster, where the devices are named 0 , 1 (same node) and ˆ0 , ˆ1 (same node), respectively.
We dispatch 128MB data with two dispatch patterns: (1) even dispatch and (2) uneven dispatchthat a greater proportion of data is exchanged with a neighbor device.
Table 1 shows the detaileddispatch proportions and the corresponding performance.
Compared with even dispatch, unevendispatch improves the overall communication performance by roughly 30%.
This is mainly becausethe 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 existingstatic even dispatch pattern is not effective enough.
There is an urgent need for a routing strategy thatcan dynamically adapt to the underlying network environments.
[Table]
Topology Aware Routing Strategy
In this section, we first abstract the data dispatch problem into an optimization objective based on thecommunication model.
Through some analysis, we obtain the target dispatch pattern under differenttopologies, which can eliminate the communication bottleneck during MoE training.
Guided by thetarget pattern, we design a topology-aware routing loss to adaptively adjust the dispatch volume.
4.1 Communication Model
We characterize the communication cost using the well-known α - β cost model, where α and βrepresent the fixed communication latency and the inverse bandwidth (i.e., transferring costs ofeach word), respectively.
For convenience, α ij and β ij are used to denote the latency and inversebandwidth between the i -th and j -th GPU.
During the training of MoE, the amount of data transferred from GPU i to E e in GPU j is c ie ∗ d ∗ b , where d ∗ b is the transferred element size.
To reduce theoverheads of multiple send-receives between two GPUs, we merge the multiple small data chunksinto a larger data chunk for delivery.
The total amount of data delivered from GPU i to GPU jis (cid:80) E ∗ ( j +
[Equation]
− 1 e =
E ∗ j c ie ∗ d ∗ b .
A global data exchange consists of P ∗ P peer-to-peer data deliveries,among which the slowest delivery, as a lower-bound, constrains the final communication performance.
Most of the global exchange implementations [20, 17, 23] are designed to approach the lower-bound.
Therefore, our ultimate objective function is to minimize the slowest send-receive communicationcost:
[Equation]
For efficient MoE training, two constraint conditions should be satisfied.
First, for any process i , thesent data size, i.e., k ∗ S , should be equal to the sum of received data size of all experts:
[Equation]
Second, to make full use of all the experts and pursue a better model accuracy, the data chunksdispatched to each expert should be balanced:
[Equation]
4.2 Model Optimization
To get the target dispatch pattern, we need to solve the optimization problem in Equation 2.
Nevertheless, Equation 2 contains plentiful parameters of a specific network, which complicates thesolving process.
Meanwhile, in some irregular topologies, some devices may suffer from quite limitedbandwidth when communicating with other devices.
According to Equation 2, the experts assignedto these devices may receive a quite small dispatch chunk size from the other processes, which maymake the experts lack of sufficient data exchanges and lead to expert isolation phenomenon.
To tacklethese problems, we simplify the optimization problem to accelerate the solving process and smooththe values of α ij , β ij for an approximate result to prevent expert isolation.
Since each send-receivecommunication shares the same α, β in homogeneous network, the target dispatch chunk size ˆ c ieis equal to the load-balanced chunk size k ∗ SN .
In the following part, we focus on the analysis of theoptimization problem under heterogeneous topologies.
On a n -layer symmetric tree topology, for any device i , all the devices can be split into n sub-groupsof G i =
{ G it | t < n } .
G it is the group of devices whose shortest path from device i are across tswitches.
Multiple hops in cross-switch communication will suffer from extra overheads and themost limited bandwidth in the hops dominates the final bandwidth.
Therefore, we can simplifythe original α ij , β ij into n value: α l =
[Equation]
) ∗
[Equation] [Equation]
2 , which can 179 precisely characterize the underlying topology and eliminate the noise of profiling.
Then we get ahierarchical matrix
ˆ α, ˆ β :
[Equation]
Take equation 3 and 4 as optimization constraint conditions, we simplify the optimization problem asfollows:
[Equation]
Then, we can get the optimal solution:
[Equation]
The optimal data distribution of the above min-max problem is only related to the bandwidth: thevolume of ˆ c ie is linear to the bandwidth.
The rationale behind the optimal result is that higherbandwidth links should bear more loads for an overall communication balanced workflow.
The ringtopology also shows a hierarchical characteristic and the solution for ring topology has the samepattern as symmetric trees.
Under some irregular asymmetric topology, the optimal result of Equation 2 may result in someexperts assigned to the devices of most limited bandwidth in lack of data exchanges with the restof the devices.
Compared to the data distribution of other experts, data chunks from local devicesoccupy a larger proportion of received data chunks in those isolated experts.
For the fairness amongthe experts, we transform the asymmetric topology into a symmetric one by merging the separatenodes into the close symmetric sub-trees.
For example, [[2,2][2]] in figure 2(d) can be merged assymmetric structure [[2,2,2]].
After that, we can optimize the lower bound of communication as thesymmetric structure.
4.3 Routing Strategy
Once getting the target data dispatch volumes among processes, we can use it to guide the MoEtraining.
To not sacrifice the model accuracy, instead of setting a compulsory dispatch ratio directly,we design a topology-aware adaptive routing loss.
[Equation]
As shown in Equation 8, we set a penalty weight p
i as the adjustment coefficients for the topology lossl itopo of each process i .
We set the normalized 1 / ˆ c i as p i to make sure c ie approximates the value ofˆ c ie .
Normalization functions that enlarge the penalty of the low-bandwidth transfer, e.g., softmax, arealso preferable.
As shown in the calculation of l itopo , the data dispatched to E e with limited bandwidthwill suffer from a larger penalty weight p ie .
Despite of these penalty modifications, auxiliary lossoccupies a small proportion of the final loss, and the value of the auxiliary loss decreases whenexperts’ number scales up.
Therefore, our final topology loss is expanded N ∗ P times to keep themagnitude of loss value.
There are a number of advantages of the proposed topology-aware strategy.
On the one hand,compared to setting a compulsory dispatch ratio, the proposed loss can adjust the communicationvolume to fit in the underlying topologies in a mild way without damaging the convergence.
Acompulsory dispatch ratio has a high potential to overwhelm the influence of the train loss andsacrifice the model accuracy.
With the topology loss, the train loss can still dominate in the finalloss value.
As a result, the dispatch results are mainly influenced by the train loss for a better modelaccuracy.
Besides, the topology-aware strategy has more potential to utilize the token information forefficient sparse training.
Guided by the topology-aware loss, the tokens nearby are more likely to beprocessed by the same expert.
Since the correlation of adjacent tokens contains the vital informationin sparse attention [10], the experts in sparsely gated routing structure may also be more likely toextract important information from adjacent tokens.
In addition, the topology-aware loss can be easily incorporated to existing MoE systems, such as
DeepSpeed-MoE and FastMoE. Taking FastMoE as an example, one just needs to directly replacethe popular load-balanced loss l aux with the proposed topology-aware loss l topo .
Since local capacitythreshold C ie is adopted in DeepSpeed-MoE for load balance, one can modify the local capacitysizes to be consistent with the proposed dispatch pattern by setting C ie to be proportional to thetarget data chunk sizes ˆ c ie .
Instead of padding the data chunks with extra zeros to be the same size as
DeepSpeed-MoE, one all-to-all communication is added to get the information of send-receive datachunk sizes and dispatch data chunks according to the sizes.
Evaluation
Experiment setup To demonstrate the effectiveness of TA-MoE, we carry out a series of experiments on three typical NVIDIA GPU clusters with different network topology, and some representative model configurations.
Table 2 lists the cluster settings.
The testbed is three typical clustersfrom the PaddleCloud 1 platform.
For cluster A, each node consists of 8 NVIDIA Tesla 40GB A
GPUs connected with NVSwitch, which shows high performance for both computation and networkcommunication.
Clusters B and C are equipped with 8 NVIDIA Tesla 32GB V100 GPUs in eachnode.
The nodes in cluster B are connected by the same switch, while cluster C is composed of a largenumber of servers and switches that are interconnected through an internal network infrastructure.
Besides, the software configurations are set as CUDA 11.0, NCCL 2.8.4 and CUDA 11.1, NCCL2.8.3, for cluster A and cluster B, C, respectively.
[Table]
Without loss of generality, we focus the experiments on problems related to language modeling, with
GPT-3 Medium [3] 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 withone expert.
Both the Switch top-1 [7] and the GShard top-2 gates [11] are tested with the weightof auxiliary loss set as 1.0.
For the consistency of the experiment, we implement the models by asingle framework Paddle [2] and train on the open-source openwebtext2 dataset [1].
More detailedspecifications of model settings can be found in Table 3.
models.
[Figure]
Methodology We incorporate TA-MoE into widely used DeepSpeed-MoE [20] and FastMoE
[8]implementations.
Because TA-MoE modifies the gate structure, we first compare the validation lossw.r.t.
steps to ensure that TA-MoE will not interfere with the convergence on various model scales.
[Figure]
On top of that, we test the overall throughput and thespeedup of TA-MoE over these two classical baselines.
To be more comprehensive, we also comparewith the recently proposed FasterMoE Hir gate [9]on the metric of time to convergence performance.
Besides, a detailed analysis of communication costs,as well as the distribution of the dispatch are alsogiven.
Accuracy and Performance We first compare thevalidation loss w.r.t steps of TA-MoE and the representative FastMoE on cluster C. As shown in Figure3, the loss curves of TA-MoE and FastMoE are consistent to converge under different training scalesof 8 experts to 48 experts.
These results demonstrate that the TA-MoE can adaptively adjust the datadispatch volume without sacrificing the model accuracy.
[Table]
[Figure]
To demonstrate the advantages of TA-MoE in terms of training performance, we compare itwith both DeepSpeed-MoE and FastMoE on various hardware and model configurations.
Theperformance indicators including throughout (tokens/s) and the speedups are depicted in Figure 4.
It is clear that TA-MoE can bring significant performance improvements over its counterparts under almost all the configurations.
When compared with the DeepSpee-MoE, theachieved speedup is about 1.05x-1.61x.
As for FastMoE, the performance improvement is around1.01x-4.77x.
More speedups are achieved on FastMoE than DeepSpeed-MoE. This is because
TA-MoE has a more dynamic dispatch pattern and larger adjustable space based on FastMoE.
[Figure]
It is also observed that more improvements are obtained on Cluster C, which can reach 4.77x for somecases, due to the relief of its serious network contention of cross-switch communication.
In addition,the comparison of the results of Switch and GShardgate reveals that TA-MoE can behave better in adjusting larger volume of data.
To be more comprehensive, we make further comparisons with the recently proposed FasterMoE [9].
Because the compulsory dispatch strategy of FasterMoE affects the convergence, we 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 5, TA-MoE can converge faster than FasterMoE.
We evaluate the time to reachthe validation loss values of 3.1, 2.9, and 2.8, and TA-MoE can converge faster by about 1.25x, 1.47xand 1.54x.
The results further verify that the proposed adaptive routing loss is more effective than thecompulsive dispatch method, which sacrifices model accuracy for training performance.
Communication Analysis To better show the effects of the dynamical dispatch strategy, we furtheranalyze communication and computation cost, and the distribution of data dispatch on cluster C. Asshown in Figure 6(a), thanks to the proposed TA-MoE strategy, the communication cost is reducedrapidly, with roughly 1.16x to 6.4x speedups.
It is also observed that the maximum speedup isachieved for 32 experts on four cross-switch nodes.
This is because the four nodes are deployedunder four different switches, and cross-switch links severely bottleneck the data exchange.
Once thetension is relieved, the obtained benefits can be dramatic.
In addition, we draw the distribution ofdata dispatch of Rank 0 as an example to observe the dispatch pattern.
In Figure 6(b), as expected,most of the data of Rank 0 are dispatched to low-overheads nearby ranks, which further verifies theeffectiveness of adaptive topology-aware loss. [Figure] u t a t i o n a n d C [Figure]
mm
[Figure]
i c a t i [Figure] t
[Figure]
[Figure] [Figure]  Conclusion In this paper, a topology-aware routing strategy, TA-MoE, was proposed to stress the mismatchbetween the data dispatch pattern and the network topology.
Based on communication modeling, weabstract the dispatch problem into an optimization objective and obtain the approximate dispatchpattern under different topologies.
On top of that, a topology-aware auxiliary loss was designed,which can adaptively route the data to fit in the underlying topology without sacrificing the modelaccuracy.
Experiments show that the proposed method can substantially outperform its counterpartson a variety of the 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.
Inthe future, we plan to take more delicate communication operators into consideration of dynamicdispatch pattern and extend our work to more hardware environments.
Checklist
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 conclusionsection
(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 conformsto them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of 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 to reproduce the mainexperimental results (either in the supplemental material or as a URL)?
[No]Please see section 4.3 and experiment setup part in section 5 for the algorithmimplementation details.
The algorithm code will be released soon.(b)
Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes] Please see the experiment setup part in section 5.
(c) Did you report error bars (e.g., with respect to the random seed after runningexperiments multiple times)?
[Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g.,type of GPUs, internal cluster, or cloud provider)?
[Yes] Please see the experimentsetup part of section 5.
If you are using existing assets (e.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 was obtained from people whose data
you’re using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identi-
fiable information or offensive content?
[N/A]
If you used crowdsourcing or conducted
research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots,
if applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
[Table]
spent on participant compensation?
[N/A]