GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding

Neural network scaling has been critical for improving the model quality in many real-world machine learning applications with vast amounts of training data and compute . Although this trend of scaling is affirmed to be a sure-fire approach for better model quality , there are challenges on the path such as the computation cost , ease of programming , and efficient implementation on parallel devices . In this paper we demonstrate conditional computation as a remedy to the above mentioned impediments , and demonstrate its efficacy and utility . We make extensive use of GShard , a module composed of a set of lightweight annotation APIs and an extension to the XLA compiler to enable large scale models with up to trillions of parameters . GShard and conditional computation enable us to scale up multilingual neural machine translation Transformer model with Sparsely-Gated Mixture-ofExperts . We demonstrate that such a giant model with 600 billion parameters can efficiently be trained on 2048 TPU v3 cores in 4 days to achieve far superior quality for translation from 100 languages to English compared to the prior art . 1 INTRODUCTION . Scaling neural networks brings dramatic quality gains over a wide array of machine learning problems such as computer vision , language understanding and neural machine translation ( Devlin et al. , 2018 ; Mahajan et al. , 2018 ; Arivazhagan et al. , 2019 ; Huang et al. , 2019 ; Brown et al. , 2020b ) . This general tendency motivated recent studies to scrutinize the factors playing a critical role in the success of scaling , including the amounts of training data , the model size , and the computation being utilized as found by past studies ( Advani & Saxe , 2017 ; Hestness et al. , 2019 ; Geiger et al. , 2020 ) . While the final model quality was found to have a power-law relationship with these factors ( Hestness et al. , 2017 ; Kaplan et al. , 2020 ) , the significant quality gains brought by larger models also came with various practical challenges . Training efficiency , which we define as the amount of compute and time used to achieve a superior model quality against the best system existed , is oftentimes left out . In this study , we strive for improving the model quality while being training efficiently . We built a 600 billion parameters sequence-to-sequence Transformer model with Sparsely-Gated Mixture-of-Experts layers , which enjoys sub-linear computation cost and O ( 1 ) compilation time . We trained this model with 2048 TPU v3 devices for 4 days on a multilingual machine translation task and achieved far superior translation quality compared to prior art when translating 100 languages to English with a single non-ensemble model . We conducted experiments with various model sizes and found that the translation quality increases as the model gets bigger , yet the total wall-time to train only increases sub-linearly with respect to the model size , as illustrated in Figure 1 . To train such an extremely large model , we relied on the following key design choices . Conditional computation First , model architecture should be designed to keep the computation and communication requirements sublinear in the model capacity . Conditional computation enables us to satisfy training and inference efficiency by having a sub-network activated on the per-input basis . Shazeer et al . ( 2017 ) has shown that scaling RNN model capacity by adding Sparsely Gated Mixture-of-Experts ( MoE ) layers allowed to achieve improved results with sub-linear cost . We therefore present our approach to extend Transformer architecture with MoE layers in this study . GShard Annotation Second , the model description should be separated from the partitioning implementation and optimization . This separation of concerns let model developers focus on the network architecture and flexibly change the partitioning strategy , while the underlying system applies semantic-preserving transformations and implements efficient parallel execution . To this end we propose a module , GShard , which only requires the user to annotate a few critical tensors in the model with partitioning policies . It consists of a set of simple APIs for annotations , and a compiler extension in XLA for automatic parallelization . Model developers write models as if there is a single device with huge memory and computation capacity , and the compiler automatically partitions the computation for the target based on the user annotations and their own heuristics . 2 MODEL . The Transformer ( Vaswani et al. , 2017 ) architecture has been widely used for natural language processing . We scale Transformer with conditional computation by replacing every other feedforward layer with a sparsely activated Position-wise Mixture of Experts ( MoE ) layer ( Shazeer et al. , 2017 ) , with a variant of top-2 gating in both the encoder and the decoder ( Figure 2 ) . Each subword token in the training example activates a sub-network of the MoE Transformer during both training and inference . The size of the sub-network is roughly independent of the number of experts per MoE Layer , allowing sublinear scaling of the computation cost . 2.1 POSITION-WISE MIXTURE-OF-EXPERTS LAYER . The Mixture-of-Experts ( MoE ) layers used in our model differ from Shazeer et al . ( 2017 ) ’ s in the sparse gating function and the auxiliary loss being used . A MoE layer for Transformer consists of E feed-forward networks FFN1 . . . FFNE , each of which outputs woe · ReLU ( wie · xs ) , where xs is the input token to the MoE layer , wi and wo being the input and output projection matrices for the feed-forward layer ( an expert ) with shapes [ M , H ] and [ H , M ] , respectively . The output of a MoE layer is the combination of the expert outputs ∑E e=1 Gs , e · FFNe ( xs ) , where the vector Gs , E is computed by a gating function GATE ( · ) . We choose to let each token dispatched to at most two experts . The corresponding gating entries Gs , e become non-zeros , representing how much an expert contributes to the final network output . The gating function GATE ( · ) is critical to the MoE layer , which is modeled by a softmax activation function to indicate the weights of each expert in processing incoming tokens . We designed a novel efficient gating function with the following mechanisms ( details illustrated in Algorithm 1 ) . Load balancing Naively picking top-k experts from the softmax probability distribution leads to load imbalance problem for training as shown in Shazeer et al . ( 2017 ) . Most tokens would have been dispatched to a small number of experts , leaving other experts insufficiently trained . To ensure the load is balanced , we enforce that the number of tokens processed by one expert is below some uniform threshold called expert capacity . Assuming N total tokens in a batch and at most two experts per token , then the expert capacity C is set to be O ( N/E ) . GATE ( · ) keeps a running counter ce for how many tokens are dispatched to an expert . When both experts selected by a token already exceed their capacity , the token is considered as an overflowed token , where Gs , E degenerates into a zero vector . Such tokens will be passed on to the next layer via residual connections . The introduction of the fixed expert capacity instead of loading balancing functions in Shazeer et al . ( 2017 ) allows us to run parallel execution of gating function as described blow . Local dispatching for parallel gating Load balancing required the token assignments of one expert dependent on assignments of the other experts . The original gating function proposed by ( Shazeer et al. , 2017 ) had to be implemented sequentially , especially under the static shape constraints on TPUs . In our study , we distributed thousands of experts over thousands of devices , a sequential implementation of the gating function would keep most of the devices idle most of the time . Instead , we propose a new GATE ( · ) function that partitions all tokens in a training batch evenly into G local groups , i.e. , each group contains S = N/G tokens for local dispatching . All local groups are processed independently in parallel . Each group is given a fractional capacity of each expert , C = 2N/ ( G · E ) , to ensure that at most this many tokens are dispatched to an expert . In general , increasing the expect capacity C decreases the number of overflowed tokens thus improves the model quality . Since G× C is a constant , however , the higher capacity leads to smaller number of groups which hurts the training throughput by limiting the number of parallel gating execution . In this way , we can ensure that expert capacity is still enforced and the overall load is balanced . With fixed expert capacity and local dispatching , we are able to speed up the gating function by O ( G ) times . Auxiliary loss Following Shazeer et al . ( 2017 ) , we define a new differentiable auxiliary loss term ` aux to enforce the load balancing . It is added to the overall loss function of the model L = ` ori +k ∗ ` aux with a constant multiplier k , where ` aux is defined in line ( 13 ) of algorithm 1 , and the term ce/S represents the fraction of input routed to each expert . We replace the mean square ( ce/S ) 2 with Algorithm 1 : Group-level top-2 gating with auxiliary loss Data : xS , a group of tokens of size S Data : C , Expert capacity allocated to this group Result : GS , E , group combine weights Result : ` aux , group auxiliary loss ( 1 ) for e← 1 to E do ( 2 ) ce ← 0 . gating decisions per expert ( 3 ) gS , e ← softmax ( wg · xS ) . gates per token per expert , wg are trainable weights ( 4 ) me ← 1S ∑S s=1 gs , e . mean gates per expert ( 5 ) end ( 6 ) for s← 1 to S do ( 7 ) g1 , e1 , g2 , e2 = top_2 ( { gs , e|e = 1 · · ·E } ) . top-2 gates and expert indices ( 8 ) g1← g1/ ( g1 + g2 ) . normalized g1 ( 9 ) c← ce1 . position in e1 expert buffer ( 10 ) if ce1 < C then ( 11 ) Gs , e1 ← g1 . e1 expert combine weight for xs ( 12 ) end ( 13 ) ce1 ← c + 1 . incrementing e1 expert decisions count ( 14 ) end ( 15 ) ` aux = 1 E ∑E e=1 ce S ·me ( 16 ) for s← 1 to S do ( 17 ) g1 , e1 , g2 , e2 = top_2 ( { gs , e|e = 1 · · ·E } ) . top-2 gates and expert indices ( 18 ) g2← g2/ ( g1 + g2 ) . normalized g2 ( 19 ) rnd← uniform ( 0 , 1 ) . dispatch to second-best expert with probability ∝ 2 · g2 ( 20 ) c← ce2 . position in e2 expert buffer ( 21 ) if c < C ∧ 2 · g2 > rnd then ( 22 ) Gs , e2 ← g2 . e2 expert combine weight for xs ( 23 ) end ( 24 ) ce2 ← c + 1 ( 25 ) end differentiable approximation me ( ce/S ) , which can provide better numerical stability since it can be optimized with gradient descent . Random routing Intuitively , the output ys is a weighted average of what selected experts return . If the weight for the 2nd expert is very small , we can simply ignore the 2nd expert to conserve the overall expert capacity . Hence , in addition to respecting the expert capacity constraint , GATE ( · ) dispatches to the 2nd-best expert with the probability proportional to its weight g2 . We observed much less overflowed tokens thus better accuracy with random routing for models at the small scale . We then adopted this approach for our experiments at large scales .

Summary: The paper applies mixture-of-experts (MoE) [1] to the Transformers to significantly increase the number of parameters in the model while keeping the total computational cost feasible. The main differences with [1] are: (1) Only choose 2 experts at each timestep; (2) Set a capacity upper bound for each expert to make sure that no expert becomes a lagger. The paper also presents a convenient library to make implementing MoE models easier. The proposed method is tested on a large multilingual machine translation dataset and shows performance gains over models trained on a single language pair and a model trained without MoE.