dPRO: A Generic Profiling and Optimization System for Expediting Distributed DNN Training

TL;DR

dPRO is a comprehensive toolkit that profiles and optimizes distributed deep neural network training across multiple frameworks, significantly improving performance by identifying bottlenecks and suggesting effective strategies.

Contribution

It introduces a unified profiling and optimization system for distributed DNN training that works across different frameworks and communication schemes, enabling performance diagnosis and acceleration.

Findings

Predicts distributed training performance with less than 5% error in most cases.

Achieves up to 3.48x speed-up over baseline methods.

Works across multiple frameworks and communication schemes.

Abstract

Distributed training using multiple devices (e.g., GPUs) has been widely adopted for learning DNN models over large datasets. However, the performance of large-scale distributed training tends to be far from linear speed-up in practice. Given the complexity of distributed systems, it is challenging to identify the root cause(s) of inefficiency and exercise effective performance optimizations when unexpected low training speed occurs. To date, there exists no software tool which diagnoses performance issues and helps expedite distributed DNN training, while the training can be run using different deep learning frameworks. This paper proposes dPRO, a toolkit that includes: (1) an efficient profiler that collects runtime traces of distributed DNN training across multiple frameworks, especially fine-grained communication traces, and constructs global data flow graphs including detailed communication operations for accurate replay; (2) an optimizer that effectively identifies performance bottlenecks and explores optimization strategies (from computation, communication, and memory aspects) for training acceleration. We implement dPRO on multiple deep learning frameworks (TensorFlow, MXNet) and representative communication schemes (AllReduce and Parameter Server). Extensive experiments show that dPRO predicts the performance of distributed training in various settings with < 5% errors in most cases and finds optimization strategies with up to 3.48x speed-up over the baselines.