Source-linked AI summary
GSPMD: General and Scalable Parallelization for ML Computation Graphs
Yuanzhong Xu, HyoukJoong Lee, Dehao Chen, Blake Hechtman, Yanping Huang, Rahul Joshi, Maxim Krikun, Dmitry Lepikhin, Andy Ly, Marcello Maggioni, Ruoming Pang, Noam Shazeer, Shibo Wang, Tao Wang, Yonghui Wu, Zhifeng Chen
TL;DR
GSPMD targets the challenge of scaling ML computations across devices without forcing users to rewrite single-device model programs or manually implement complete sharding plans. It infers graph-wide partitioning from a few tensor annotations, supports mixed parallelism patterns, and demonstrates scaling across diverse models and thousands of TPU cores, with homogeneous pipeline stages as a scope boundary.
Problem
Neural-network scaling creates a need to parallelize both training data and model parameters across devices.
Method
GSPMD uses tensor-sharding annotations and compiler passes to infer and implement graph-wide partitioning, including combinations of parallelism paradigms.
Results
GSPMD achieves close-to-linear memory and performance scaling across image, speech, sparse-language, and dense-language models as device count increases.
Takeaways & Limitations
GSPMD provides a reusable parallelization mechanism for scaling models across domains and up to thousands of Cloud TPUv3 cores.
Takeaways & Limitations
The pipeline reduction is limited to homogeneous pipeline stages; heterogeneous stages require integration with other pipeline implementations.
Abstract
from arXiv · showhide
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.
1 Introduction
GSPMD addresses the need to scale increasingly large neural networks by using simple tensor-sharding annotations to unify and automate multiple parallelism strategies. It supports production-oriented compilation and partitioning constraints while demonstrating scaling across model types and thousands of TPU cores.
- Neural-network scaling creates demand for parallelizing both training data and model parameters through data, pipeline, and within-layer model parallelism.
- GSPMD unifies data, in-layer model, image-spatial, and weight-update or optimizer-state sharding, while supporting combinations across and within layers.A wrapper library can reduce pipeline parallelism to tensor and operator partitioning in common cases.
- GSPMD generalizes beyond mixture-of-experts models and has helped scale language, image, and speech models.
- GSPMD uses a few tensor-distribution annotations, then completes sharding across the computation graph and generates an equivalent parallelized computation.Users can change partitioning strategies by reconfiguring annotations rather than rewriting model code.
- GSPMD produces one SPMD program for all partitions, supports uneven dimensions on arbitrary device meshes, and covers XLA operators through its compiler integration.These choices target compilation scalability and compatibility with statically known accelerator shapes.
- GSPMD achieves close-to-linear memory and performance scaling with device count across image, speech, sparse-language, and dense-language model use cases.
2 Background
Modern ML workloads use several parallelism patterns, and GSPMD provides a general in-operator partitioning framework that can combine them. Pipeline parallelism requires a wrapper in supported cases, while GSPMD’s compiler passes operate on XLA’s graph representation.
- Data parallelism replicates the model across devices, processes different training data, and synchronizes local gradients with AllReduce.
- Within-layer model parallelism shards layer weights across devices, potentially requiring communication, and can enable larger models by distributing weights.
- Spatial partitioning shards image inputs along spatial dimensions to fit large image data on devices with limited memory.
- Weight-update or optimizer-state sharding distributes gradient application across replicas, especially benefiting expensive optimizers such as ADAM.
- Pipeline parallelism divides the training graph into stages that exchange forward and backward results, with dependency-induced idle-time bubbles.Larger global training batches can amortize bubble overhead.
- GSPMD natively supports in-operator parallelism, can reduce pipeline parallelism to tensor partitioning with a wrapper, and expresses nested combinations of techniques.
- The automatic partitioner is implemented as XLA compiler transformation passes over HLO dataflow graphs, reducing implementation burden across front-end frameworks.
3 Tensor Sharding and Auto Completion
GSPMD represents tensor sharding with a small set of general abstractions and automatically completes operator shardings from user annotations. Compatible shardings support mixed data, model, optimizer-state, expert, and pipeline parallelism.
- Sharding representation: GSPMD separates sharding completion from per-operator partitioning, assigning every tensor an explicit or inferred property describing its device distribution.Its representation supports replicated, tiled, and partially tiled shardings.
- Sharding representation: The representation treats all tensor dimensions uniformly rather than specializing for batch, weight, or image dimensions.This lets one abstraction express data, model, weight-update, and spatial parallelism.
- Sharding representation: mesh_split maps tensor dimensions to device-mesh dimensions, expressing tiled, partially tiled, or replicated sharding depending on how many mesh dimensions are mapped.Users can configure mesh order for communication topology and may use different meshes for different tensors.
- Sharding completion: Additional weight partitioning can trigger weight-update sharding, reducing peak memory through on-demand unsharding and using ReduceScatter for gradients.Optimizer updates are applied on a sharded optimizer.
- Pipeline parallelism: GSPMD reduces constrained pipeline parallelism to tensor sharding by shifting a state buffer, with extra iterations representing pipeline bubbles.Annotating the stage dimension lets GSPMD lower buffer shifting to CollectivePermute, while the implementation also combines with other parallelism patterns.
- Sharding completion: Compatible shardings from different operands can be merged, allowing simple annotations to combine batch and model dimensions in the result.For a fully connected layer, batch sharding expresses data parallelism while feature sharding expresses model parallelism.
4 The SPMD Partitioner
The SPMD partitioner addresses production challenges by compiling one program across partitions while handling static shapes, communication, windowed operations, and nested sharding patterns.
- SPMD implementation: GSPMD chooses SPMD over MPMD to avoid prohibitively slow compilation of many partition-specific programs at large partition counts.A single program works for all partitions, while MPMD creates a customized program for each partition.
- Static shapes: Uneven partitions remain compatible with static shapes by rounding dimensions up to partition-count multiples and handling padded regions.Padded data may require masking so it does not affect operator results.
- Communication: GSPMD uses collective communication, including AllGather, AllReduce, AllToAll, CollectivePermute, and halo exchange, to implement regular cross-partition patterns.Halo exchange uses CollectivePermute to exchange partial data with neighboring partitions.
- Windowed operators: Windowed operators may require halo exchange together with padding, slicing, and masking because neighboring partitions need overlapping data and configurations can be uneven.GSPMD supports convolution configurations such as arbitrary padding and dilation, using maximum halo sizes and DynamicSlice when needed.
- Grouping and recursive partitioning: Recursive pattern matching handles rank-polymorphic operators and nested sharding cases without manually writing rules for every dimension configuration.This enables combining spatial and feature partitioning within a single Convolution operator.
- Grouping and recursive partitioning: Partition grouping creates nested contexts that reinterpret logical partition IDs as subgroups of physical devices for recursive partitioning.In the illustrated Einsum case, the inner partitioner creates an AllGather after rewriting logical groups according to the context.
5 Case Study and Evaluation
GSPMD is evaluated across dense, sparse, hybrid, speech, and image models, using tensor sharding to combine parallelism strategies and scale training across TPU devices. The experiments show near-linear memory or performance scaling in several settings, while pipelining introduces measurable overhead.
- Dense Transformer language model: GSPMD trains Transformer models with hundreds of billions of parameters using annotations on only 7 tensors per layer, while inferring shardings for internal computations and normalization.Those annotations represent roughly 0.7% of all tensors in the XLA graph.
- Dense Transformer language model: A 2D mesh shards model dimensions across X and Y to fit large weights while preserving TPU compute efficiency.The initial configuration shards H and N along Y and M along X, but partially sharded activations and small per-device weights limit scalability and efficiency.
- Combining pipelining and in-layer sharding: 24% slower than 2D sharding, the best pipelined configuration combines 4 stages with 4 in-layer model-parallel shards but incurs bubble and recomputation overheads.Rematerialization reduces peak memory and permits more microbatches, helping amortize pipeline bubbles.
- Sparse Transformer language model: 0.98s to 1.10s step time is maintained as sparse MoE experts grow from 32 to 512, but 2048 experts raise step time to 1.51s.At the largest scale, AllToAll communication and gating computation become more significant.
- Image spatial partitioning: 15.7x step-time reduction on 16 partitions demonstrates nearly linear spatial-partitioning scaling for a 128x128x128 image, while 32-way partitioning doubles FLOPS utilization in a larger image setting.The higher utilization results from enabling a larger per-device batch size.
6 Related Work
GSPMD extends distributed ML parallelization beyond basic data and graph partitioning by providing a general sharding-based transformation mechanism. It supports diverse and combined strategies while complementing automated policy search.
- Frameworks and generality: GSPMD expands beyond framework-level data parallelism and graph-partitioned model parallelism through SPMD-style per-operator partitioning.The related-work comparison positions GSPMD as a mechanism for transforming annotated computation graphs.
- Frameworks and generality: GSPMD supports broader partitioning than Tofu, including all dimensions of complex operators such as convolution.Tofu is described as supporting limited strategies, whereas GSPMD handles complex operator dimensions.
- Pipeline parallelism: GSPMD can express pipelining through vectorization or combine pipeline stages with additional partitioning within each stage.Other pipeline scheduling techniques remain orthogonal and can reduce pipeline bubbles.
- Memory and optimizer-state sharding: GSPMD provides a uniform annotation API for partitioning weights, activations, and optimizer state without distinguishing tensor categories.The cited comparison contrasts this generality with Zero’s separate optimizations for different tensor types.
- Combined parallelism: GSPMD implements many prior combined model-parallel and pipelining techniques, including automatic scatter/gather behavior across pipeline stages.The cited example uses fully sharded activations for scattering and on-demand combination in the next stage.
- Automated partitioning: GSPMD and FlexFlow are complementary: automated search can define a partitioning space, while GSPMD transforms annotated graphs.FlexFlow focuses on discovering partitioning policies; GSPMD focuses on graph-transformation mechanisms.
7 Conclusion
GSPMD offers largely automated, annotation-driven parallelization for machine-learning computations. The paper reports broad model coverage across image, speech, and language workloads with scalable performance and memory behavior on thousands of TPU cores.
- Conclusion: GSPMD partitions image, speech, and language models across up to thousands of Cloud TPUv3 cores.The conclusion reports good and predictable performance and memory scaling across these models.
- Conclusion: GSPMD combines a simple API for multiple parallelism patterns with auto-completion that requires annotations on only a few tensors.The system completes the partitioning of the entire model from limited user guidance.
A.1 XLA operators for dynamism
GSPMD uses XLA operators to support dynamic behaviors across partitions despite XLA’s requirement for static shapes.
- XLA operators for dynamism: Although XLA requires static shapes, GSPMD supports data access with dynamic offsets based on run-time values.Table 9 summarizes operators used for dynamic behavior across partitions.
- XLA operators for dynamism: Table 9 lists XLA operators that GSPMD uses to handle non-uniform behavior between partitions.The table is presented as a summary of operators supporting dynamic behavior across partitions.
A.2 Halo exchange details
GSPMD handles partitioned convolutions by accounting for stride, padding, dilation, and non-uniform halos. Its halo-exchange procedures use padding, slicing, masking, and case-specific alignment to preserve valid windows across partitions.
- Window configurations: Stride is the distance between successive window positions, while low/high padding adds elements at the two ends of the base dimension.These window parameters describe how convolution inputs and output positions are arranged.
- Window configurations: Base dilation inserts elements between base values before low/high padding, whereas window dilation inserts elements between window values.Base and window dilation affect different operands and are handled differently during partitioning.
- Non-constant halo exchange: Non-constant halo sizes can vary by partition, such as right halos of 1, 2, 3, and 4 for a four-way convolution partition.The example expresses halo size as partition_id + 1.
- Non-constant halo exchange: General halo exchange performs maximum-size exchange, partition-specific DynamicSlice, and masking of invalid or garbage values.The sequence first exchanges the largest halos, then selects each partition’s valid region and masks out-of-range data.
- Base dilation: Base dilation complicates halo exchange because partition offsets may fall in dilation holes, producing different edge and interior behavior.GSPMD handles this through three cases based on the relationship among stride, per-shard window count, and dilation.
- Base dilation: When stride × per-shard window count is divisible by dilation, partitions share the same initial padding configuration and exchange halos on the undilated base region.The divisibility condition aligns partition starts before the first base data element.
- Base dilation: For stride 1 with misaligned per-shard window counts, maximum low padding can be used because every padded and dilated base position is a valid window start.The passage contrasts this case with applying Pad and DynamicSlice before dilation, which would scale padding incorrectly.
- Base dilation: When non-unit stride causes dilation misalignment, additional window padding masks unaligned base elements so valid window starts remain aligned.This addresses cases where choosing any single low-padding amount would make some partition invalid.