# GSPMD: General and Scalable Parallelization for ML Computation Graphs

**Authors:** 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  
**arXiv:** [2105.04663](https://arxiv.org/abs/2105.04663)  
**Date:** 2021-05-10  
**Summary page:** https://www.summarizepaper.com/en/paper/2105.04663v1/

## Key points

- General and Scalable Parallelization for ML Computation Graphs (GSPMD)
- Innovative system designed to optimize model performance in distributed computing environments
- Compiler-based approach allows writing programs similar to single-device ones
- Efficiently parallelizes computation process by leveraging user annotations on tensor distribution preferences
- Accommodates various paradigms of parallelism through simple yet versatile representation of partitioning
- Impressive Compute Utilization Rates
- Demonstrated rates ranging from 50% to 62% on configurations utilizing between 128 to 2048 Cloud TPUv3 cores for models with up to one trillion parameters
- Unified Program Generation
- Generates a unified program catering to all devices while dynamically adjusting behavior according to run-time partition ID
- Leverages collective operators for efficient cross-device communication purposes
- Scalability and Consistent Compilation Time
- Scalability extends beyond computational performance as compilation time remains consistent even with an increasing number of devices in use
- Significant Advancement in Parallelization Systems
- Offers both generalizability and scalability, representing a significant advancement in parallelization systems for machine learning computation graphs

## Summary

The <Organization>GSPMD</Organization>: General and Scalable Parallelization for ML Computation Graphs presents an innovative system designed to optimize model performance in distributed computing environments. Developed by authors 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 and Zhifeng Chen,<Organization>GSPMD</Organization> offers a compiler-based approach that allows users to write programs similar to single-device ones. By leveraging user annotations on tensor distribution preferences within the computation graph structure,<Organization>GSPMD</Organization> efficiently parallelizes the computation process. Its simple yet versatile representation of partitioning accommodates various paradigms of parallelism and can automatically infer partitioning for each operator based on limited user annotations. This feature streamlines the scaling up of existing single-device programs by addressing technical challenges such as static shape constraints and uneven partitioning issues. Through advanced techniques and optimizations,<Organization>GSPMD</Organization> has demonstrated impressive compute utilization rates ranging from 50% to 62% on configurations utilizing between 128 to 2048 Cloud TPUv3 cores for models with up to one trillion parameters. Additionally,<Organization>GSPMD</Organization> generates a unified program that caters to all devices while dynamically adjusting its behavior according to a run-time partition ID. The system also leverages collective operators for efficient cross-device communication purposes. Notably,<Organization>GSPMD</Organization>'s scalability extends beyond computational performance as the compilation time remains consistent even with an increasing number of devices in use. Overall, <Organization>GSPMD</Organization> represents a significant advancement in parallelization systems for machine learning computation graphs by offering both generalizability and scalability.