GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding

The problem

Serving translation between many languages can be done with a separate model per language pair. The GShard paper measures itself against exactly that arrangement: translating from 100 languages into English would otherwise mean training 100 bilingual baseline models ^[1]. Replacing that crowd of models with a single multilingual system means the system needs far more capacity, which pushes the parameter count into the hundreds of billions. At that size a model no longer fits comfortably on one accelerator, and the work of splitting it across many devices by hand becomes the real bottleneck — every change to the model has to be matched by a change to the splitting code.

The paper is an attempt to remove that bottleneck: to keep the model code recognisable while still expressing computations spread over thousands of chips.

How it works

GShard is described as a module composed of a set of lightweight annotation APIs plus an extension to the XLA compiler. Together they provide a way to express a wide range of parallel computation patterns with minimal changes to existing model code ^[1]. The idea is that the model author marks up the parts of the network that should be split across devices, and the compiler works out how to partition the tensors and insert the communication that the split requires — the paper's term for this is automatic sharding.

The second ingredient is conditional computation, in the form of a sparsely gated mixture of experts. In a mixture-of-experts layer, several expert subnetworks sit side by side and a gate decides which experts each input is sent to, so that only part of the network is used for any one token. Parameter count and computation per token are thereby decoupled: the model can be very large while each token is still processed by a small fraction of it. The authors used this mechanism to scale a multilingual neural machine translation Transformer beyond 600 billion parameters through automatic sharding ^[1].

Put the two halves together and the division of labour is clear. The mixture of experts supplies capacity without a proportional rise in compute, and the annotation-and-compiler layer makes the resulting giant model expressible and trainable without hand-written partitioning code.

Results

The giant model was efficiently trained on 2048 TPU v3 accelerators in 4 days ^[1]. On translation from 100 languages to English it achieved far superior quality compared with the prior art ^[1].

The cost accounting is as much a result as the quality numbers. The authors put the training cost of the 600B model at 22 TPU v3 core-years, against 29 TPU v3 core-years for training all 100 bilingual baseline models ^[1]. One very large sparse model was therefore not only better on quality but cheaper to train than the collection of smaller dense models it replaced.

Why it matters

Large-model work is often described in terms of parameter counts, but parameter counts on their own say little about whether a model can actually be trained. The GShard paper makes two contributions that speak to that gap. First, it shows a concrete route from sparse conditional computation to a working model at a scale of hundreds of billions of parameters, and reports the compute cost openly, so the trade-off between size and training budget can be judged rather than guessed. Second, it treats the distribution of a model across accelerators as a compiler problem rather than a rewriting problem: because the sharding is expressed through annotations, a wide range of parallel patterns can be described while the model code stays close to its original form ^[1].

Trade-offs that remain

Sparse routing buys capacity, but it also places demands of its own on the system. Tokens are sent to different experts, which may live on different devices, so training involves communication between chips that dense models do not need; the balance of load across experts has to be managed, or some devices idle while others work; and the benefit of using only part of the network per token depends on the hardware and interconnect underneath. The paper's contribution is a way to express and train such models at scale, and the pattern it sets out — conditional computation plus automatic sharding — leaves those engineering questions to whoever builds on it.