Zorse: Optimizing LLM Training Efficiency on Heterogeneous GPU Clusters

TL;DR

Zorse is a system designed to improve the efficiency of training large language models on heterogeneous GPU clusters by integrating flexible parallelism strategies and an automatic configuration planner.

Contribution

It introduces Zorse, the first system to unify pipeline and data parallelism with adaptive configuration for heterogeneous GPU clusters.

Findings

Zorse outperforms existing systems in heterogeneous training scenarios.

It effectively balances load and memory across diverse GPUs.

The automatic planner optimizes training strategies for specific workloads.

Abstract

Large language models (LLMs) require vast amounts of GPU compute to train, but limited availability and high costs of GPUs make homogeneous clusters impractical for many organizations. Instead, assembling heterogeneous clusters by pooling together GPUs of different generations allows them to achieve higher aggregate compute and make use of all available GPUs. However, training on heterogeneous clusters presents several challenges, including load balancing across GPUs, optimizing memory usage to accommodate varying memory capacities, and ensuring communication-efficient training over diverse network interconnects potentially spanning multiple datacenters. In this paper, we make the case that efficient training on heterogeneous clusters requires (1) the integration of pipeline parallelism and data parallelism in a manner that is both communication- and memory-efficient, and (2) a more adaptable configuration of pipeline and data parallelism, which includes the capability to flexibly partition GPUs into asymmetric pipeline parallel stages and to incorporate heterogeneous GPUs within the same data parallelism group. We propose Zorse, the first system to unify all these capabilities while incorporating a planner that automatically configures training strategies for a given workload. Our evaluation shows that Zorse significantly outperforms state-of-the-art systems in heterogeneous training scenarios.