GSPMD: General and Scalable Parallelization for ML Computation Graphs (2105.04663v2)

Abstract: 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.

• The paper introduces GSPMD, a compiler-based system that automatically partitions ML computation graphs to support diverse parallelism paradigms including data, model, and pipeline parallelism.
• It achieves significant compute utilization, reaching 50% to 62% on up to 2048 Cloud TPUv3 cores while scaling models up to one trillion parameters.
• The paper presents an SPMD execution model with recursive and nested partitioning that simplifies sharding and minimizes manual configuration in distributed ML training.

Analyzing GSPMD: A Compiler-Based Approach for Scalable ML Parallelization

The paper presents GSPMD, a compiler-based system for parallelizing machine learning models across distributed devices. GSPMD is designed to support the automatic partitioning of computation graphs with minimal user input, allowing models to scale efficiently on high-performance hardware like Cloud TPUv3 cores. The system allows users to program machine learning models as though they are to be run on a single device and then specify tensor distribution through a few annotations, which the GSPMD then uses to implement parallelization.

Key Features and Results

1. Unified Representation of Parallelism: GSPMD is capable of expressing various parallelism paradigms, such as data parallelism, in-layer model parallelism, and pipeline parallelism. This is achieved through a flexible sharding API that utilizes a logical device mesh to determine how tensors are distributed. This unified approach simplifies the programming model by abstracting parallelism concerns away from the core model logic.
2. Scalability and Efficiency: GSPMD achieves significant compute utilization when running large scale models, reaching 50% to 62% utilization on upto 2048 Cloud TPUv3 cores. By supporting various partition strategies and combining them effectively, GSPMD supports models with up to one trillion parameters.
3. Automatic Sharding Completion: The system includes mechanisms to infer the partitioning strategy for an entire graph based on user annotations made to a few tensors. This reduces the burden on the user to devise complex sharding schemes while allowing for intuitive interactions with the overall partitioning strategy. The sharding completion includes propagation through graph nodes based on operator semantics, merging compatible strategies, and iterative refinement to reach optimal configurations without user intervention.
4. Single Program Multiple Data (SPMD) Execution: GSPMD utilizes an SPMD approach, creating a single program suitable for all partitions instead of individual ones for each partition. This reduces compilation time and complexity and is vital in achieving scalability across thousands of partitions.
5. Recursive and Nested Partitioning: The system introduces a recursive framework for partitioning operators with rank polymorphism, such as Einsum and Convolution. This framework enables partitioning that captures nested dimensions efficiently, allowing for the expression of mixed parallelism patterns without extensive manual configuration.
6. Handling Complex Production Environments: GSPMD addresses practical challenges such as static shape constraints, where tensor shapes must be known at compile-time. Techniques like halo exchange, dynamic bounds, and handling of padding are used to manage these constraints within SPMD partitioning effectively.

Implications and Future Directions

Practically, GSPMD simplifies the adaptation of existing models written for single-device execution to run on large distributed systems, facilitating model scaling without in-depth parallelization expertise. Theoretically, it advances the state of compiler-assisted parallelization by generalizing and expanding previous techniques into a robust, scalable system.

Looking ahead, GSPMD's integration with AI frameworks such as TensorFlow and its compatibility with hardware backends like TPUs position it well to benefit from advancements in multi-device functionalities and AI compiler optimizations. Future research could explore automated search methods integrated with GSPMD for optimizing partition strategies in more complex pipelines beyond its current semi-automatic approach.

The paper positions GSPMD as a versatile tool in the distributed ML training landscape, and likely future developments could leverage its foundations to integrate more adaptive and automated parallelism discovery. The success of GSPMD in real-world scenarios underscores its potential for broader application across diverse machine learning domains.