GSPMD: General and Scalable Parallelization for ML Computation Graphs

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.

References in corpus (4)

• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• Language Models are Few-Shot Learners
• TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models
• Efficient Algorithms for Device Placement of DNN Graph Operators

Cited by in corpus (6)

• BigSSL: Exploring the Frontier of Large-Scale Semi-Supervised Learning for Automatic Speech Recognition
• Improving Automatic Parallel Training via Balanced Memory Workload Optimization
• Amazon SageMaker Model Parallelism: A General and Flexible Framework for Large Model Training
• Skillful Twelve Hour Precipitation Forecasts using Large Context Neural Networks
• Memory-efficient array redistribution through portable collective communication
• DistIR: An Intermediate Representation and Simulator for Efficient Neural Network Distribution