Expediting Distributed DNN Training With Device Topology-Aware Graph Deployment

Abstract

This paper presents TAG, an automatic system to derive optimized DNN training graph and its deployment onto any device topology, for expedited training in device-and topology-heterogeneous ML clusters. We novelly combine both the DNN computation graph and the device topology graph as input to a graph neural network (GNN), and join the GNN with a search-based method to quickly identify optimized distributed training strategies. To reduce communication in a heterogeneous cluster, we further explore a lossless gradient compression technique and solve a combinatorial optimization problem to automatically apply the technique for training time minimization. We evaluate TAG with various representative DNN models and device topologies, showing that it can achieve up to 4.56x training speed-up as compared to existing schemes. TAG can produce efficient deployment strategies for both unseen DNN models and unseen device topologies, without heavy fine-tuning.

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

“…(3) DeepSpeed [36] supports ZeRO-based [35] data parallelism and implements intra-op model parallelism for MoE layers. (4) TAG [56] is a heterogeneity-aware DNN training system. TAG supports data parallelism and inter-op model parallelism.…”

“…The performance of SFB is primarily determined by the batch size and the number of devices involved [7,54]. TAG [56] proposes an integer linear programming-based technique to automatically identify beneficial application of SFB to tensors in a DNN model trained in a homogeneous cluster. However, uneven tensor partitioning across heterogeneous resources introduces additional complication to this problem.…”

“…On the other hand, the rapid evolution of accelerate devices (e.g., GPUs, TPUs, Habana chips) and multi-tenant resource sharing have resulted in mixed device types and uneven interconnect bandwidth in many clusters. Efficient exploitation of available heterogeneous resources for DNN training has piqued significant interest from AI practitioners [16,25,27,30,43,53,56]. A prevalent approach for making use of heterogeneous clusters is to distribute multiple training jobs onto different homogeneous subsets of the cluster [50,57].…”

“…The limited existing heterogeneity-aware DNN training designs mostly support data parallelism and inter-op model parallelism [27,53,56]. For data parallelism, each device trains the DNN model for a distinct portion of the dataset, allocated according to capacity of the device, and synchronizes model gradients among each other at the end of each training iteration.…”