PaperlayerView exact revision on arXiv

Source-linked AI summary

GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding

Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, Zhifeng Chen

arXiv:2006.16668v1cs.CLcs.LGstat.ML

TL;DR

Scaling neural networks improves quality but raises challenges in computation cost, programming effort, and parallel implementation. GShard combines lightweight sharding annotations, automatic compiler partitioning, and sparsely gated Mixture-of-Experts to scale multilingual Transformers. The resulting 600B-parameter model was trained on 2048 TPU v3 accelerators in 4 days and achieved superior translation quality for 100 languages to English.

  • Problem

    Scaling giant neural networks requires efficient computation, programming abstractions, and parallel-device implementations beyond the capacity of conventional model-parallel systems.

  • Method

    GShard uses lightweight sharding annotations and an XLA compiler extension for automatic parallelization, combined with conditional computation through sparsely gated Mixture-of-Experts Transformers.

  • Results

    A 600B-parameter multilingual Transformer was trained on 2048 TPU v3 accelerators in 4 days and achieved superior translation quality for 100 languages to English.

  • Takeaways & Limitations

    Conditional computation and automatic sharding provide a favorable trade-off between model scale and computational cost while simplifying large-scale parallel execution.

  • Takeaways & Limitations

    The compared models have disproportionate training resource requirements, and parameter count alone does not always reflect effective capacity in massively multitask settings.

Abstract

from arXiv · show

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.

1 Introduction

GShard addresses the practical barriers to scaling neural networks by combining conditional computation with lightweight sharding annotations and compiler-based automatic parallelization. The approach supports a 600B-parameter multilingual Transformer trained efficiently across 2048 TPU v3 devices.

  • 1 Introduction: Scaling improves model quality but creates challenges in training efficiency, model parallelism, infrastructure scalability, and manual partitioning.These challenges involve compute and time costs, under-utilized devices, large distributed graphs, and substantial engineering effort.
  • 1 Introduction: A 600B-parameter multilingual Transformer was trained on 2048 TPU v3 devices for 4 days while achieving superior translation quality across 100 languages.The model uses Sparsely-Gated Mixture-of-Experts layers and has sublinear computation cost and O(1) compilation time.
  • 1 Introduction: GShard uses conditional computation to keep computation and communication requirements sublinear as model capacity grows.Sparsely activated sub-networks allow capacity to increase without proportionally increasing per-input computation.
  • 1 Introduction: GShard separates model description from partitioning through lightweight tensor annotations and an XLA compiler extension for automatic parallelization.Developers describe models as if they ran on one large device, while the compiler partitions computation using annotations and heuristics.
  • 1 Introduction: SPMD partitioning generates one program for all devices, avoiding the scalability problems of separate MPMD programs.For the illustrated dot product, devices compute local results and combine them with an AllReduce.

2 Model

The MoE Transformer replaces selected feed-forward layers with sparsely gated expert networks, so each token activates only a small sub-network. Its gating design must balance expert load while remaining efficient for very large batches and expert counts.

  • Transformer layers combine self-attention with position-wise feed-forward computation, while the decoder additionally uses cross-attention.
  • Each token activates a sub-network whose size is roughly independent of the number of experts, enabling sublinear computation scaling.
  • The MoE layer uses multiple feed-forward experts and a gating network that dispatches each token to at most two experts.
  • Replacing every other feed-forward layer with an MoE layer produces the MoE Transformer Encoder structure.
  • Top-k routing can overload a few experts, so gating must balance load and avoid sequential O(N E) computation when N and E are very large.The passage describes N in the order of millions and E in the order of thousands.
  • Expert capacity limits each expert to a uniform token threshold, set to O(N/E) when each token is dispatched to at most two experts.

3 Highly Parallel Implementation using GShard

GShard combines tensor sharding annotations with compiler transformations to express and execute large-scale parallel computations across thousands of devices. Its implementation handles MoE computation, communication, uneven partitions, and padding while preserving a logical full-shape program.

  • Compiler transformation: The compiler converts partially annotated computations into SPMD programs, inserts cross-device communication, and handles irregular partitioning patterns.It generates one program launched on all devices, keeping compilation time near constant as partitions increase.
  • MoE implementation: GShard expresses MoE execution with tensor operations including dispatch, expert computation, and combination through Einsum operations.Top-2 gating sends each token to at most two experts, while dispatch and combine tensors encode expert destinations and buffer positions.
  • Scaling assumptions: The implementation analyzes complexity under assumptions including G = O(D), E = O(D), and constant per-device tokens, hidden size, and expert dimensions.These assumptions support the stated scaling analysis and include practical device-memory and capacity constraints.
  • Sharding APIs: GShard provides replicate, split, and shard APIs for annotating tensor replication, partitioning, and device placement.Annotations separate model description from parallel implementation and do not change logical tensor shapes.
  • Flexible partitioning: The APIs support sharding batch, feature, expert, and spatial dimensions, allowing different model components to use different partitioning strategies.This flexibility supports switching partition modes between MoE and non-MoE layers.
  • Irregular partitions: Uneven partitions are padded and may require masking with an operator’s identity value, such as zero for Reduce-Add, to preserve correctness.The masking procedure uses partition offsets and full-shape offsets to select valid operand values or the identity value.

4 Massively Multilingual, Massive Machine Translation (M4)

Massively multilingual translation balances positive transfer for low-resource languages against capacity bottlenecks and negative transfer. GShard’s sparsely gated MoE scaling improves quality and training efficiency across model depths, expert counts, and resource regimes.

  • Motivation: A single multilingual model must improve both high-resource translation through added capacity and low-resource translation through positive transfer.Increasing the number of language-pair tasks can deliver larger gains for low-resource languages, while model scaling addresses capacity constraints.
  • Approach: GShard combines conditional computation with sparsely gated MoE Transformers to scale beyond 1 trillion parameters while keeping training practical.A 600B-parameter GShard model processes 1T tokens in 250k steps in under 4 days; each token is routed through a small number of experts.
  • Depth and transfer: Increasing depth from 12L to 36L yields 2-to-3 BLEU points on average for both low- and high-resource languages when experts per layer remain fixed.Deeper models provide consistent upward shifts in ΔBLEU, while increasing experts primarily benefits high-resource languages.
  • Capacity scaling: +3.3 average BLEU across 100 languages results when experts per layer increase from 128 to 512, but scaling from 512 to 2048 adds only +1.3 average BLEU.The results indicate diminishing returns after the capacity bottleneck is substantially relaxed.
  • Depth and transfer: Dense-deep T(96L) transfers better to low-resource languages, while MoE(36E, 128L) achieves the same transfer quality with 37 billion parameters.The paper attributes the dense model’s advantage to greater shared subnetworks, while more experts relax the capacity bottleneck but reduce shared transfer.
  • Training efficiency: MoE models with three times the depth require 2 to 3 times fewer tokens to reach preset training-loss thresholds.For example, MoE(128E, 12L) takes 3 times as many tokens as MoE(128E, 36L) to reach 0.7 training cross-entropy.
  • Training efficiency: 600 billion parameters trained on 2048 TPU cores for 4 days achieves the best average BLEU, while dense T(96L) requires 235 TPU core years and trails GShard models in quality.The 600B MoE model itself requires 22.4 TPU years, motivating continued attention to cost-effective scaling.

5 Performance and Memory Consumption

GShard maintains roughly constant per-device memory as experts scale and achieves sublinear runtime growth through partitioned computation and communication. Measurements identify AllToAll dispatch as the dominant scaling cost, while many operators remain efficient across thousands of devices.

  • Overall scalability: 1.7x execution time increase accompanies a 16x model scale from 128 to 2048 devices, while device memory consumption remains roughly constant.The measurements cover computation and memory efficiency on TPUs as both device and expert counts increase.
  • Memory scaling: O(1) per-device memory scaling holds as experts increase because replicated weights, distributed MoE weights, and activations have constant per-device sizes after SPMD partitioning.With fixed layers, both weight and activation memory stay constant as the number of experts increases.
  • Runtime efficiency: At 128 experts, the model reaches >70% of roofline performance, while 2048 experts increase device time by 1.7x and retain 48% of roofline performance.The roofline is an optimistic estimate based on peak compute, memory-bandwidth, or interconnect-bandwidth utilization.
  • Operator performance: Dense Transformer computations achieve >85% peak FLOPS and maintain constant per-device cost as the number of experts increases.These operations are mainly large matrix multiplications that use the TPU matrix unit effectively.
  • Operator performance: Gating overhead is limited by small-constant-factor O(D) computations, with Cumsum taking less than 10% of total MoE and Transformer layer time at 2048 experts.The Cumsum operations are memory-bound or sequential, but their execution time remains relatively small at the measured scale.
  • Communication performance: AllToAll execution time grows about 3.75x when experts increase 16x from 128 to 2048, raising its runtime share from 16% to 36%.AllToAll implements MoE dispatch and combine, and its cost scales roughly as O(√D).
  • Communication performance: AllToAll increases only 9x when partitions grow 128x from 16 to 2048, supporting efficient cross-partition dispatch through resharding.AllReduce execution time is independent of device count, whereas AllToAll becomes more expensive sublinearly.
  • Operator scalability: Across partitioned operators, most compute and communication scale sublinearly, while spatially partitioned convolutions exhibit O(1) scaling.The operator measurements are intended to guide applications beyond the MoE model.

6 Related Work

Related work establishes that larger neural networks often improve performance and describes the hardware, software, parallelism, and conditional-computation foundations for scaling them. Prior approaches motivate automated parallelism and sparsely gated mixture-of-experts Transformers.

  • Model scaling: Larger models have empirically improved performance across vision, language understanding, translation, and other complex tasks.Examples include deeper ResNets, larger Transformers, and models with larger vocabularies or embeddings.
  • Hardware and software: Accelerator hardware and software libraries jointly support neural-network computation, but accelerators remain difficult to program directly.Frameworks abstract hardware-specific details while relying on lower-level libraries to drive accelerators efficiently.
  • Parallel training: Data parallelism replicates a program across devices with different inputs, whereas model parallelism partitions computation or parameters across devices.Modern neural-network training and inference commonly use clusters containing multiple accelerators.
  • Automated parallelism: Automated parallelism systems aim to reduce the burden of specifying distributed computation, including graph partitioning and SPMD-style programming.Existing frameworks and systems provide varying levels of support for data parallelism, graph partitioning, and operator-level partitioning.
  • Conditional computation: Mixture-of-experts models have been used to increase model capacity while activating only input-dependent sub-networks.Conditional computation routes examples according to criteria such as difficulty, computation budget, or learned sparse expert selection.

7 Conclusion

GShard combines automatic computation partitioning with conditional computation to make giant neural-network scaling practical and efficient. The study reports quality gains, favorable training-cost trade-offs, and practical lessons about separating model design from parallelization.

  • GShard scaled a multilingual MoE Transformer to 600B parameters, training it in 4 days for translation from 100 languages to English.The model achieved superior translation quality compared with prior art.
  • 22 TPU v3 core years versus 29 TPU years for 100 bilingual Transformer baselines demonstrates lower reported training cost for the MoE models.
  • Progressive scaling produced consistent quality gains, suggesting that quality improvements had not plateaued at the studied scales.
  • Conditional computation offered a favorable trade-off between model scale and computational cost, enabling trillion-parameter experiments in days rather than weeks or months.
  • Separating model description from parallelization lets developers focus on network implementation while GShard partitions computation graphs and generates programs for all devices.
  • Parameter count alone does not always track effective model capacity at scale, especially in massively multi-task settings with imbalanced training data.

A.1 Decoding with Flat Beam Search

The decoding procedure uses autoregressive beam search, requiring repeated decoder execution and cross-device communication in each decoder MoE layer. Flattening beam hypotheses avoids repeated key/value reordering, though its efficiency depends on implementation details.

  • Autoregressive decoding executes the decoder stack m times for an output of length m, with decoder MoE layers requiring cross-device dispatch and combine operations.
  • Flattening beam hypotheses into one interleaved sequence and modifying the attention mask lets the decoder avoid reordering previously computed key/value tensors after each beam expansion.
  • The flattening trade-off can be positive or negative because memory bandwidth limits affect incremental Transformer decoding.

A.2 Machine Translation Experiments Details

The machine-translation experiments use a shared Transformer configuration, Adafactor optimization, and multilingual SentencePiece tokenization. These details define the model dimensions, optimization schedule, and source–target vocabularies used in the experiments.

  • MoE Transformer models share a 1024-dimensional Transformer representation, 8192-dimensional feed-forward and MoE hidden layers, 16 attention heads, and 128-dimensional keys and values.Input, residual, and attention dropout are all 0.1.
  • Adafactor uses β1 = 0.0, β2 = 0.99 with a 1 −t−0.8 schedule, update clipping at 1.0, and square-root learning-rate decay after 10k steps.The initial learning rate is 1.0.
  • SentencePiece provides a 64,000-item multilingual source vocabulary spanning 102 languages and a 32,000-item English-only target vocabulary.

A.3 General Sharding API

GShard exposes lightweight sharding APIs that annotate tensors and support device-aware partitioning. Device assignments encode how tensor slices map to devices, while topology-aware layouts can reduce costly data movement.

  • Advanced sharding strategies can minimize data transfers beyond the replicate() and split() APIs.
  • shard(tensor, device_assignment) annotates a tensor for partitioning according to a multidimensional array of device IDs.The assignment has the tensor's rank, and its elements identify the devices occupying corresponding data slices.
  • A 2D tensor can be divided into 2x4 partitions, with communication occurring between partitions along tensor rows; device numbers identify the assigned devices.
  • Device assignments should account for target topology and inter-partition communication because data movement strongly affects parallel-execution performance.Figure 10 compares assignments based on topology and row-wise communication patterns.

A.4 SPMD Partitioning for Convolution and Window-Based Operators

GShard partitions convolution and other window-based operators across spatial dimensions, propagating sharding through subsequent layers and the backward pass. Its SPMD partitioner handles halo exchange, including partition-dependent halos and base dilation.

  • Spatial partitioning: GShard partitions convolutional spatial dimensions and propagates the resulting sharding to other layers and the backward pass.The same framework supports window-based operations such as ReduceWindow.
  • Halo exchange: Window-based operators require halo exchange because data may be shared between neighboring windows.GShard uses CollectivePermute, which requires statically shaped communication, creating complications when halo sizes differ.
  • Window configurations: The partitioner considers stride, low/high padding, base dilation, and window dilation when determining convolution behavior.Base dilation applies to the left-hand-side input, while window dilation applies to the right-hand-side kernel.
  • Halo exchange: Non-constant halo sizes are handled by exchanging the maximum halo across partitions, slicing each partition's valid region, and removing garbage values.The general sequence is described as maximum-size exchange, partition-specific DynamicSlice, and cleanup of out-of-range data.
  • Base dilation: Base dilation requires three cases because partitions can begin at dilation holes and padding behaves differently at edges than in interior regions.The cases depend on stride and whether stride × per_shard_window_count is divisible by dilation.
  • Window configurations: Window dilation is simpler than base dilation because the partitioned right-hand side has no low/high padding.The discussion notes that this commonly arises during gradient computation for strided convolutions.
Loading 2006.16668v1…