Diving into 3D Parallelism with Heterogeneous Spot Instance GPUs: Design and Implications

TL;DR

This paper analyzes the challenges of 3D parallelism in heterogeneous GPU environments for large language model training, and introduces AutoHet, a system that optimizes parallelism plans and supports elastic training with spot instances.

Contribution

AutoHet is a novel system that automatically optimizes 3D parallelism strategies for heterogeneous GPUs and enables efficient elastic training with preemption recovery.

Findings

AutoHet achieves up to 1.79× training throughput speedup.

AutoHet reduces recovery time by 4.38× compared to baseline.

Theoretical model effectively balances load across diverse GPUs.

Abstract

The rapid growth of large language models (LLMs) and the continuous release of new GPU products have significantly increased the demand for distributed training across heterogeneous GPU environments. In this paper, we present a comprehensive analysis of the challenges involved in implementing 3D parallelism in such environments, addressing critical issues such as the need for symmetric tensor parallelism, efficient gradient synchronization in asymmetric pipeline parallelism, and the trade-offs between memory utilization and computational efficiency. Building upon these insights, we introduce AutoHet, a novel system that automatically identifies the optimal parallelism plan for distributed training on heterogeneous GPUs. AutoHet supports asymmetric 3D parallelism structures and facilitates fine-grained workload distribution. We propose a theoretical model that frames the device grouping and load balancing as an optimization problem to minimize per-iteration training time, thus effectively balancing computing power and memory usage across GPUs with diverse capabilities. To enable elastic training upon spot instance preemption, AutoHet presents an efficient recovery strategy that prioritizes to retrieve training states from local nodes, and only downloads the missing checkpoints from the cloud storage. Our extensive evaluation, conducted on three large-scale models and utilizing combinations of three different GPU types, demonstrates that AutoHet outperforms existing DNN training systems, achieving up to a 1.79$\times$ speedup in training throughput compared with Megatron-LM and Whale, and a 4.38$\times$ speedup of recovery speed compared to a spot instance baseline.