mlsanityData to 7 Oct 2026

GSPMD: General and Scalable Parallelization for ML Computation Graphs

arXiv 2105.04663 ↗First posted 10 May 2021cs.DC cs.LGcited by 109

Not on the map. Why?

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.

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

Builds on

21 of 44 references.

References · 21

  1. 2017Attention Is All You Needcited by 59,932
  2. 2014Adam: A Method for Stochastic Optimizationcited by 36,960
  3. 2020Language Models are Few-Shot Learnerscited by 32,199
  4. 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,160
  5. 2021Zero-Shot Text-to-Image Generationcited by 4,276
  6. 2017Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layercited by 3,413
  7. 2020Conformer: Convolution-augmented Transformer for Speech Recognitioncited by 1,838
  8. 2019Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelismcited by 1,829
  9. 2020GShard: Scaling Giant Models with Conditional Computation and Automatic Shardingcited by 1,601
  10. 20163D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotationcited by 1,471
  11. 2018Adafactor: Adaptive Learning Rates with Sublinear Memory Costcited by 896
  12. 2016Training Deep Nets with Sublinear Memory Costcited by 825
9 more
  1. 2018Beyond Data and Model Parallelism for Deep Neural Networkscited by 177
  2. 2018Mesh-TensorFlow: Deep Learning for Supercomputerscited by 168
  3. 2020DAPPLE: A Pipelined Data Parallel Approach for Training Large Modelscited by 100
  4. 2021TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Modelscited by 74
  5. 2018Glow: Graph Lowering Compiler Techniques for Neural Networkscited by 73
  6. 2018Supporting Very Large Models using Automatic Dataflow Graph Partitioningcited by 43
  7. 2020Automatic Cross-Replica Sharding of Weight Update in Data-Parallel Trainingcited by 30
  8. 2018Relay: A New IR for Machine Learning Frameworkscited by 24
  9. 2020Efficient Algorithms for Device Placement of DNN Graph Operatorscited by 24

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2021: 11; 2022: 17; 2023: 25; 2024: 19; 2025: 21; 2026: 16.

109 papers.

Sort:

By area:

  1. 2022PaLM: Scaling Language Modeling with Pathwayscited by 5,304
  2. 2024Gemma 2: Improving Open Language Models at a Practical Sizecited by 1,928
  3. 2025Gemma 3 Technical Reportcited by 1,686
  4. 2022Scaling Autoregressive Models for Content-Rich Text-to-Image Generationcited by 1,287
  5. 2022CoCa: Contrastive Captioners are Image-Text Foundation Modelscited by 1,265
  6. 2022LaMDA: Language Models for Dialog Applicationscited by 1,218
  7. 2023PaLM 2 Technical Reportcited by 1,154
  8. 2021GLaM: Efficient Scaling of Language Models with Mixture-of-Expertscited by 890
  9. 2024Gemma: Open Models Based on Gemini Research and Technologycited by 889
  10. 2023PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallelcited by 641
  11. 2022Mixture-of-Experts with Expert Choice Routingcited by 611
  12. 2022Efficiently Scaling Transformer Inferencecited by 529
38 more
  1. 2022Reducing Activation Recomputation in Large Transformer Modelscited by 300
  2. 2025MAGI-1: Autoregressive Video Generation at Scalecited by 262
  3. 2023Challenges and Applications of Large Language Modelscited by 251
  4. 2022Unified Scaling Laws for Routed Language Modelscited by 249
  5. 2023Google USM: Scaling Automatic Speech Recognition Beyond 100 Languagescited by 228
  6. 2023TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddingscited by 198
  7. 2022Scaling Up Models and Data with textttt5x and textttseqiocited by 187
  8. 2026Gemma 4 Technical Reportcited by 187
  9. 2024OpenMoE: An Early Effort on Open Mixture-of-Experts Language Modelscited by 185
  10. 2021Combined Scaling for Zero-shot Transfer Learningcited by 182
  11. 2023Noise2Music: Text-conditioned Music Generation with Diffusion Modelscited by 173
  12. 2021Sequence Parallelism: Long Sequence Training from System Perspectivecited by 125
  13. 2025Cambrian-S: Towards Spatial Supersensing in Videocited by 123
  14. 2021Large-Scale Differentially Private BERTcited by 117
  15. 2023Skeleton-of-Thought: Prompting LLMs for Efficient Parallel Generationcited by 103
  16. 2024DistriFusion: Distributed Parallel Inference for High-Resolution Diffusion Modelscited by 98
  17. 2021BigSSL: Exploring the Frontier of Large-Scale Semi-Supervised Learning for Automatic Speech Recognitioncited by 88
  18. 2022Decentralized Training of Foundation Models in Heterogeneous Environmentscited by 87
  19. 2022Zeus: Understanding and Optimizing GPU Energy Consumption of DNN Trainingcited by 78
  20. 2023Deep Learning for Day Forecasts from Sparse Observationscited by 70
  21. 2022When does dough become a bagel? Analyzing the remaining mistakes on ImageNetcited by 60
  22. 2022Galvatron: Efficient Transformer Training over Multiple GPUs Using Automatic Parallelismcited by 51
  23. 2024On Scaling Up 3D Gaussian Splatting Trainingcited by 49
  24. 2023AlpaServe: Statistical Multiplexing with Model Parallelism for Deep Learning Servingcited by 31
  25. 2025MegaScale-Infer: Serving Mixture-of-Experts at Scale with Disaggregated Expert Parallelismcited by 31
  26. 2023The Efficiency Spectrum of Large Language Models: An Algorithmic Surveycited by 25
  27. 2024AsyncDiff: Parallelizing Diffusion Models by Asynchronous Denoisingcited by 22
  28. 2021Skillful Twelve Hour Precipitation Forecasts using Large Context Neural Networkscited by 20
  29. 2022Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learningcited by 20
  30. 2022On Optimizing the Communication of Model Parallelismcited by 19
  31. 2023IRGen: Generative Modeling for Image Retrievalcited by 18
  32. 2024ORBIT: Oak Ridge Base Foundation Model for Earth System Predictabilitycited by 17
  33. 2023MAD Max Beyond Single-Node: Enabling Large Machine Learning Model Acceleration on Distributed Systemscited by 14
  34. 2023Improving Automatic Parallel Training via Balanced Memory Workload Optimizationcited by 13
  35. 2026Scaling Mixture-of-Experts Video Pretraining for Embodied Intelligencecited by 13
  36. 2021Automap: Towards Ergonomic Automated Parallelism for ML Modelscited by 12
  37. 2023Blockwise Parallel Transformer for Large Context Modelscited by 12
  38. 2021Amazon SageMaker Model Parallelism: A General and Flexible Framework for Large Model Trainingcited by 11

50 most cited of 109 All 109 (JSON)