# GSPMD: General and Scalable Parallelization for ML Computation Graphs

arXiv: https://arxiv.org/abs/2105.04663 · HTML: https://mlsanity.com/paper/2105.04663/

## Key facts

- Cited by 111 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 44 references that are in the map.
- First posted on arXiv 10 May 2021.

## Builds on (most cited first, 21 of 21)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Conformer: Convolution-augmented Transformer for Speech Recognition](https://mlsanity.com/paper/2005.08100/) (2020)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation](https://mlsanity.com/paper/1606.06650/) (2016)
- [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://mlsanity.com/paper/1804.04235/) (2018)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [Beyond Data and Model Parallelism for Deep Neural Networks](https://mlsanity.com/paper/1807.05358/) (2018)
- [Mesh-TensorFlow: Deep Learning for Supercomputers](https://mlsanity.com/paper/1811.02084/) (2018)
- [DAPPLE: A Pipelined Data Parallel Approach for Training Large Models](https://mlsanity.com/paper/2007.01045/) (2020)
- [Glow: Graph Lowering Compiler Techniques for Neural Networks](https://mlsanity.com/paper/1805.00907/) (2018)
- [TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models](https://mlsanity.com/paper/2102.07988/) (2021)
- [Supporting Very Large Models using Automatic Dataflow Graph Partitioning](https://mlsanity.com/paper/1807.08887/) (2018)
- [Automatic Cross-Replica Sharding of Weight Update in Data-Parallel Training](https://mlsanity.com/paper/2004.13336/) (2020)
- [Relay: A New IR for Machine Learning Frameworks](https://mlsanity.com/paper/1810.00952/) (2018)
- [Efficient Algorithms for Device Placement of DNN Graph Operators](https://mlsanity.com/paper/2006.16423/) (2020)

## Built on by (most cited first, 30 of 111)

- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [Gemma 2: Improving Open Language Models at a Practical Size](https://mlsanity.com/paper/2408.00118/) (2024)
- [Gemma 3 Technical Report](https://mlsanity.com/paper/2503.19786/) (2025)
- [Scaling Autoregressive Models for Content-Rich Text-to-Image Generation](https://mlsanity.com/paper/2206.10789/) (2022)
- [CoCa: Contrastive Captioners are Image-Text Foundation Models](https://mlsanity.com/paper/2205.01917/) (2022)
- [LaMDA: Language Models for Dialog Applications](https://mlsanity.com/paper/2201.08239/) (2022)
- [PaLM 2 Technical Report](https://mlsanity.com/paper/2305.10403/) (2023)
- [Gemma: Open Models Based on Gemini Research and Technology](https://mlsanity.com/paper/2403.08295/) (2024)
- [GLaM: Efficient Scaling of Language Models with Mixture-of-Experts](https://mlsanity.com/paper/2112.06905/) (2021)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [Mixture-of-Experts with Expert Choice Routing](https://mlsanity.com/paper/2202.09368/) (2022)
- [Efficiently Scaling Transformer Inference](https://mlsanity.com/paper/2211.05102/) (2022)
- [Reducing Activation Recomputation in Large Transformer Models](https://mlsanity.com/paper/2205.05198/) (2022)
- [MAGI-1: Autoregressive Video Generation at Scale](https://mlsanity.com/paper/2505.13211/) (2025)
- [Gemma 4 Technical Report](https://mlsanity.com/paper/2607.02770/) (2026)
- [Challenges and Applications of Large Language Models](https://mlsanity.com/paper/2307.10169/) (2023)
- [Unified Scaling Laws for Routed Language Models](https://mlsanity.com/paper/2202.01169/) (2022)
- [Google USM: Scaling Automatic Speech Recognition Beyond 100 Languages](https://mlsanity.com/paper/2303.01037/) (2023)
- [TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings](https://mlsanity.com/paper/2304.01433/) (2023)
- [Scaling Up Models and Data with textttt5x and textttseqio](https://mlsanity.com/paper/2203.17189/) (2022)
- [OpenMoE: An Early Effort on Open Mixture-of-Experts Language Models](https://mlsanity.com/paper/2402.01739/) (2024)
- [Combined Scaling for Zero-shot Transfer Learning](https://mlsanity.com/paper/2111.10050/) (2021)
- [Noise2Music: Text-conditioned Music Generation with Diffusion Models](https://mlsanity.com/paper/2302.03917/) (2023)
- [Cambrian-S: Towards Spatial Supersensing in Video](https://mlsanity.com/paper/2511.04670/) (2025)
- [Sequence Parallelism: Long Sequence Training from System Perspective](https://mlsanity.com/paper/2105.13120/) (2021)
- [Large-Scale Differentially Private BERT](https://mlsanity.com/paper/2108.01624/) (2021)
- [Skeleton-of-Thought: Prompting LLMs for Efficient Parallel Generation](https://mlsanity.com/paper/2307.15337/) (2023)
- [DistriFusion: Distributed Parallel Inference for High-Resolution Diffusion Models](https://mlsanity.com/paper/2402.19481/) (2024)
- [BigSSL: Exploring the Frontier of Large-Scale Semi-Supervised Learning for Automatic Speech Recognition](https://mlsanity.com/paper/2109.13226/) (2021)
- [Decentralized Training of Foundation Models in Heterogeneous Environments](https://mlsanity.com/paper/2206.01288/) (2022)

## Abstract (arXiv)

We present GSPMD, an automatic, compiler-based parallelization system for common machine learning computations. It allows users to write programs in the same way as for a single device, then give hints through a few annotations on how to distribute tensors, based on which GSPMD will parallelize the computation. Its representation of partitioning is simple yet general, allowing it to express different or mixed paradigms of parallelism on a wide variety of models. GSPMD infers the partitioning for every operator based on limited user annotations, making it convenient to scale existing single-device programs. It solves several technical challenges for production usage, allowing GSPMD to achieve 50% to 62% compute utilization on up to 2048 Cloud TPUv3 cores for models with up to one trillion parameters.