Galvatron: An Automatic Distributed System for Efficient Foundation Model Training
Xinyi Liu, Yujie Wang, Shenhan Zhu, Fangcheng Fu, Qingshuo Liu, Guangming Lin, Bin Cui
TL;DR
The paper tackles the challenge of efficiently training large-scale Foundation Models under diverse workloads and hardware by introducing Galvatron, an automatic distributed training system. It combines a hardware-and-model profiler, a cost-model driven search engine, and a runtime that supports data, tensor, pipeline, sharded data, and sequence parallelism along with recomputation to automatically compose a cost-efficient hybrid strategy at per-layer granularity. Key contributions include automatic parallelism optimization via decision-tree modeling and dynamic programming, end-to-end automation from profiling to execution, and open-source availability with documentation. Empirical results on clusters with NVIDIA GPUs show throughput improvements of up to $1.26\times$–$1.47\times$ over Megatron and DeepSpeed, with adoption in industry underscoring practical impact.
Abstract
Galvatron is a distributed system for efficiently training large-scale Foundation Models. It overcomes the complexities of selecting optimal parallelism strategies by automatically identifying the most efficient hybrid strategy, incorporating data, tensor, pipeline, sharded data, and sequence parallelism, along with recomputation. The system's architecture includes a profiler for hardware and model analysis, a search engine for strategy optimization using decision trees and dynamic programming, and a runtime for executing these strategies efficiently. Benchmarking on various clusters demonstrates Galvatron's superior throughput compared to existing frameworks. This open-source system offers user-friendly interfaces and comprehensive documentation, making complex distributed training accessible and efficient. The source code of Galvatron is available at https://github.com/PKU-DAIR/Hetu-Galvatron.