HAP: SPMD DNN Training on Heterogeneous GPU Clusters with Automated Program Synthesis

TL;DR

HAP tackles efficient SPMD DNN training on heterogeneous GPU clusters by automatically synthesizing a distributed program that emulates a single-device TensorFlow/PyTorch graph. It alternates between optimizing the distributed program $Q$ and tensor sharding ratios $B$, using syntax-guided program synthesis with an A*-based search and a linear-programming load balancer, while integrating heterogeneity-aware communication strategies such as Padded All-Gather/Grouped Broadcast and SFB. The system is implemented in PyTorch with a Rust-based synthesizer and demonstrates up to $2.41\times$ speedups on heterogeneous clusters and strong performance on homogeneous clusters, with synthesis overhead of only seconds for sizable models. By jointly optimizing computation, communication, and load balancing, HAP enables scalable training of large models (including MoE variants) on mixed hardware, and is released to the community for broader experimentation and adoption.

Abstract

Single-Program-Multiple-Data (SPMD) parallelism has recently been adopted to train large deep neural networks (DNNs). Few studies have explored its applicability on heterogeneous clusters, to fully exploit available resources for large model learning. This paper presents \OurSystem, an automated system designed to expedite SPMD DNN training on heterogeneous clusters. \OurSystem jointly optimizes the tensor sharding strategy, sharding ratios across heterogeneous devices and the communication methods for tensor exchanges for optimized distributed training with SPMD parallelism. We novelly formulate model partitioning as a program synthesis problem, in which we generate a distributed program from scratch on a distributed instruction set that semantically resembles the program designed for a single device, and systematically explore the solution space with an A*-based search algorithm. We derive the optimal tensor sharding ratios by formulating it as a linear programming problem. Additionally, \OurSystem explores tensor communication optimization in a heterogeneous cluster and integrates it as part of the program synthesis process, for automatically choosing optimal collective communication primitives and applying sufficient factor broadcasting technique. Extensive experiments on representative workloads demonstrate that \OurSystem achieves up to 2.41x speed-up on heterogeneous clusters.

Figures (19)

• Figure 1: Common collective communication operations.
• Figure 2: Performance with different sharding ratios under different computation-to-communication ratios (computed for P100 GPUs).
• Figure 3: Different implementations of All-Gather for uneven shards.
• Figure 4: Performance of padded All-Gather and grouped Broadcast under different tensor sharding ratios.
• Figure 5: Sufficient Factor Broadcasting. (b) and (c) depict the SPMD program on each device.