GSPMD: General and Scalable Parallelization for ML Computation Graphs

By Yuanzhong Xu, HyoukJoong Lee, Dehao Chen, Blake A. Hechtman, Yanping Huang, Rahul Joshi, Maxim Krikun, Dmitry Lepikhin, Andy Ly, Marcello Maggioni and 5 more. 2021

188 citations recorded by the source corpus.

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.

Subfields: Computer Networks and Communications, Computer Vision and Pattern Recognition, Hardware and Architecture, 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 Computer Networks and Communications

In Computer Vision and Pattern Recognition

In Hardware and Architecture

In Efficiency & Compression

Citations in the corpus

PaLM: Scaling Language Modeling with Pathways (2022, methodology): “Each TPU v4 Pod contains a full copy of the model parameters, with each weight tensor partitioned over 3072 chips using 12-way model parallelism and 256-way fully sharded data parallelism (the approach termed “2D finalized” in Xu et al. (2021)).”
Unified Scaling Laws for Routed Language Models (2022, methodology): “We note there are several libraries available for supporting large-scale Routing, including DeepSpeed [Kim et al., 2021, Rajbhandari et al., 2022] and GSPMD [Xu et al., 2021].”
LaMDA: Language Models for Dialog Applications (2022, methodology): “The largest model in LaMDA was pre-trained on 1024 TPU-V3 chips and 123 TFLOPS/s for 57.7 days with FLOPS utilization of 56.5% using GSPMD [95].”
GLaM: Efficient Scaling of Language Models with Mixture-of-Experts (2021, methodology): “We rely on GSPMD’s compiler pass (Xu et al., 2021) to automatically determine the sharding properties for the rest of the tensors.”
Scaling Up Models and Data with t5x and seqio (2022, methodology): “While XLA GSPMD will automatically select matching partitions for the intermediate activations produced by these parameters, users may also provide overrides by naming their axes as well via flax.partitioning.with sharding constraint to better optimize memory usage and between-device communication.”
PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel (2023, methodology): “Further, compiler-based techniques such as Alpa [35], GSPMD [31], and FlexFlow [9] leverage profiling, performance modeling, user annotations and search to find the best configuration across the parallelism space of data, tensor and pipeline for a given cluster.”
Efficiently Scaling Transformer Inference (2022, neutral): “Prior works propose several approaches for efﬁcient partitioning to train large models efﬁ-ciently, for e.g., NeMo Megatron (Korthikanti et al., 2022), GSPMD (Xu et al., 2021) and Alpa (Zheng et al., 2022).”
Mixture-of-Experts with Expert Choice Routing (2022, methodology): “The large MoE model is partitioned with a 2D sharding algorithm as presented in GSPMD [36], which fully exploits the 2D topology of the TPU cluster [19].”
BigSSL: Exploring the Frontier of Large-Scale Semi-Supervised Learning for Automatic Speech Recognition (2021, methodology): “We use the GShard [62] framework with the GSPMD backend [63] to train the 8B model on Cloud TPUs.”
Combined Scaling for Zero-shot Transfer Learning (2021, methodology): “Our approach is based on the Single-Program Multiple-Data (SPMD) technique, which has been successfully applied to train large language models in previous works such as in Xu et al. (2021); Lepikhin et al. (2020).”
Scaling Autoregressive Models for Content-Rich Text-to-Image Generation (2022, methodology): “There are two additional benefits in adopting GSPMD pipelining: 1) it allows us to conveniently configure pipelines within model sub-components, simplifying the overall complexity for encoderdecoder models, and 2) since pipelining is implemented as tensor partitioning on vectorized programs, we can ”
CoCa: Contrastive Captioners are Image-Text Foundation Models (2022, methodology): “Our models are implemented in the Lingvo framework [45] with GSPMD [46, 47, 48, 49] for scaling performance.”
Combined Scaling for Open-Vocabulary Image Classiﬁcation (2022, methodology): “Our approach is based on the Single-Program Multiple-Data (SPMD) technique, which has been successfully applied to train large language models in previous works such as in Xu et al. (2021); Lepikhin et al. (2020).”
This paper is included in the Proceedings of the 16th USENIX Symposium on Operating Systems Design and Implementation. (unknown): “Mesh-USENIX TensorFlow [48], GSPMD [31,57] and OneFlow [58] provide annotation APIs for users to manualy speciﬁy the intra-op parallel plan.”
Overlap Communication with Dependent Computation via Decomposition in Large Deep Learning Models (2022, methodology): “Intra-layer (tensor) model parallelism [45, 46, 52] is one of the techniques to enable large models.”
Gemma 2: Improving Open Language Models at a Practical Size (2024, methodology): “As in Gemma 1, we use the GSPMD partitioner (Xu et al., 2021) for training step computation and the MegaScale XLA compiler (XLA, 2019).”
TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings (2023, background): “The “1D/2D activation/weight partitioning” option is a typical way of partitioning tensors in a large model graph (see Figure 7 in [63]).”