Improving Automatic Parallel Training via Balanced Memory Workload Optimization

TL;DR

This paper introduces Galvatron-BMW, a system that automatically finds the most efficient hybrid parallelism strategies for training Transformer models across multiple GPUs, improving throughput and resource utilization.

Contribution

The paper presents a novel framework that automates hybrid parallelism strategy selection using decision trees and dynamic programming, optimizing training efficiency.

Findings

Galvatron-BMW outperforms previous methods in training throughput.

It effectively balances workload across GPUs under memory constraints.

The system adapts to different Transformer models and hardware setups.

Abstract

Transformer models have emerged as the leading approach for achieving state-of-the-art performance across various application domains, serving as the foundation for advanced large-scale deep learning (DL) models. However, efficiently training these models across multiple GPUs remains a complex challenge due to the abundance of parallelism options. Existing DL systems either require manual efforts to design distributed training plans or limit parallelism combinations to a constrained search space. In this paper, we present Galvatron-BMW, a novel system framework that integrates multiple prevalent parallelism dimensions and automatically identifies the most efficient hybrid parallelism strategy. To effectively navigate this vast search space, we employ a decision tree approach for decomposition and pruning based on intuitive insights. We further utilize a dynamic programming search algorithm to derive the optimal plan. Moreover, to improve resource utilization and enhance system efficiency, we propose a bi-objective optimization workflow that focuses on workload balance. Our evaluations on different Transformer models demonstrate the capabilities of Galvatron-BMW in automating distributed training under varying GPU memory constraints. Across all tested scenarios, Galvatron-BMW consistently achieves superior system throughput, surpassing previous approaches that rely on limited parallelism strategies.