Parmesan: Efficient Partitioning and Mapping Flow for DNN Training on General Device Topology

Recently, various pipeline parallelism strategies are proposed to tackle the scalability problem of training a large DNN model on a distributed system. However, most of the works focus on pipeline scheduling while lacking a general methodology to handle network partitioning and mapping to distributed systems with heterogeneous interconnection. In this work, we propose an efficient design flow, named Parmesan, to map the training of a large DNN onto a system with general device topology to maximize the throughput. Parmesan works in an end-to-end manner and solves the whole optimization problem in two phases. The first phase aims at producing well-balanced partitions, and the second phase works towards placing the DNN on devices connected by an arbitrary topology network, considering the heterogeneity of the interconnection bandwidth. We show that Parmesan speeds up the pipeline training throughput on systems with different GPU topologies and is able to handle the mapping problem for heterogeneously interconnected architectures. We believe our proposed general device topology mapping algorithm will provide valuable information for architecture designers and assist them in designing a more DNN-friendly architecture.