GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding

By Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, Z. Chen. 2020

2,182 citations recorded by the source corpus.

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. GShard is a module composed of a set of lightweight annotation APIs and an extension to the XLA compiler. It provides an elegant way to express a wide range of parallel computation patterns with minimal changes to the existing model code. GShard enabled us to scale up multilingual neural machine translation Transformer model with Sparsely-Gated Mixture-of-Experts beyond 600 billion parameters using automatic sharding. We demonstrate that such a giant model can efficiently be trained on 2048 TPU v3 accelerators in 4 days to achieve far superior quality for translation from 100 languages to English compared to the prior art.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Large Language Models, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Large Language Models

In Efficiency & Compression

Citations in the corpus

PaLM: Scaling Language Modeling with Pathways (2022, background): “Many other works aim to increase of the scale of models, while limiting communication overheads (Rajbhandari et al., 2020; Lepikhin et al., 2020; Li et al., 2020; Rasley et al., 2020; Rajbhandari et al., 2021; Ren et al., 2021; Narayanan et al., 2021a).”
Mixtral of Experts (2024, methodology): “Taking this all together, the output y for an input token x is computed as: This formulation is similar to the GShard architecture [21], with the exceptions that we replace all FFN sub-blocks by MoE layers while GShard replaces every other block, and that GShard uses a more elaborate gating…”
An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale (2020, neutral): “Thanks to Transformers’ computational efﬁciency and scalability, it has become possible to train models of unprecedented size, with over 100B parameters (Brown et al., 2020; Lepikhin et al., 2020).”
Unified Scaling Laws for Routed Language Models (2022, methodology): “Routing was popularized for large scale training by Shazeer et al. [2017], and furthered by work including GShard [Lepikhin et al., 2020], Switch Transformer [Fedus et al., 2021] and GLaM [Du et al., 2021].”
Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model (2022, methodology): “Some prior works have chosen to use the mixture-of-experts (MoE) [32, 35, 61] technique to scale to larger model sizes more economically, producing large-scale models that selectively use a subset of its parameters in each forward pass.”
Efficient Large Scale Language Modeling with Mixtures of Experts (2021, methodology): “In Table 7 we report training time and energy usage for our models based on the above estimates, as well as estimates from Patterson et al. (2021) for other large-scale LMs, in particular GShard (Lepikhin et al., 2021), Switch Transformer (Fedus et al., 2021) and GPT-3 (Brown et al., 2020).”
GLaM: Efficient Scaling of Language Models with Mixture-of-Experts (2021, methodology): “In particular, Lepikhin et al. (2021) showed that a 600B parameter MoE model achieved +6.1 average BLEU across 100 language pairs compared to the best dense prior-art with 10.7× saving in computation cost.”
Pathways: Asynchronous Distributed Dataflow for ML (2022, methodology): “Models like Mixture of Experts (MoE) (Shazeer et al., 2017) and routed capsule networks (Hinton et al., 2018; Barham and Isard, 2019) exploit computational sparsity by “routing” different (sub-)examples to the accelerators hosting different subsets of model weights based on learned functions that ar”
Finetuned Language Models Are Zero-Shot Learners (2021, background): “In addition to the traditional “dense” models, sparse mixture of experts (MoE) models of up to more than 1T parameters have been trained and finetuned (Lepikhin et al., 2020; Fedus et al., 2021).”
GSPMD: General and Scalable Parallelization for ML Computation Graphs (2021, background): “• Partially tiled (an extension to GShard [23]). organized in a logical multi-dimensional tensor called a device mesh , which can be used with an intuitive API. mesh_split(tensor, device_mesh, dims_mapping) is sharding annotation for tensor , based on the device mesh i dims_mapping[i] -1 to…”
PanGu-α: Large-scale Autoregressive Pretrained Chinese Language Models with Auto-parallel Computation (2021, methodology): “One regime is op-level model parallelism [17, 18, 19, 20, 21, 22, 23], which partitions its involved tensors of each operator (layer), as shown in Figure 4(b).”
Carbon Emissions and Large Neural Network Training (2021, background): “…mixture-of-expert-style models can provide more than 10X reductions in computation requirements and energy costs for both training and inference while providing significantly higher accuracy than dense Transformer or LSTM-based models of equivalent computational cost per token [Sha17,Lep20,Fed21].”
Beyond English-Centric Multilingual Machine Translation (2020, background): “To this end, previous MMT work has considered different types of large capacity models (Arivazhagan et al., 2019; Lepikhin et al., 2020).”
Emergent Abilities of Large Language Models (2022, neutral): “One direction is using sparse mixture-of-experts architectures (Lepikhin et al., 2021; Fedus et al., 2021; Artetxe et al., 2021), which scale up the number of parameters in a model while maintaining constant computational costs for an input.”
Efficiently Scaling Transformer Inference (2022, neutral): “Sparsity techniques, such as task-based mixture of expert architectures (Fedus et al., 2022; Kudugunta et al., 2021; Lepikhin et al., 2020; Shazeer et al., 2017), and adaptive computation techniques that allocate different amounts of compute per input and generation timestep (Jaszczur et al., 2021;…”
Fast Inference of Mixture-of-Experts Language Models with Offloading (2023, background): “Since then, several lines of work explore Mixture-of-Experts with Transformer-based language models for machine translation Lepikhin et al. (2020) 2020), but require more parameters than a dense model with equivalent perplexity.”
A Review of Sparse Expert Models in Deep Learning (2022, methodology): “Many earlier works scaled to thousands of relatively small experts per layer, which has produced excellent pre-training and translation quality (Shazeer et al., 2017; Lepikhin et al., 2020; Fedus et al., 2021).”
Mixture-of-Experts with Expert Choice Routing (2022, methodology): “This distribution denotes the preference over experts of each incoming token, which is computed similarly in a conventional gating network [10, 21, 31].”