Beyond Data and Model Parallelism for Deep Neural Networks

TL;DR

This paper introduces FlexFlow, a framework that searches a comprehensive space of parallelization strategies for deep neural network training, significantly improving throughput and scalability over existing methods.

Contribution

The paper defines a new, broader search space for parallelization strategies called SOAP and develops FlexFlow, a system that efficiently finds optimal strategies using a fast simulation-based search.

Findings

FlexFlow achieves up to 3.8x training throughput improvement.

The SOAP search space enables more effective parallelization strategies.

FlexFlow scales better across multiple GPUs.

Abstract

The computational requirements for training deep neural networks (DNNs) have grown to the point that it is now standard practice to parallelize training. Existing deep learning systems commonly use data or model parallelism, but unfortunately, these strategies often result in suboptimal parallelization performance. In this paper, we define a more comprehensive search space of parallelization strategies for DNNs called SOAP, which includes strategies to parallelize a DNN in the Sample, Operation, Attribute, and Parameter dimensions. We also propose FlexFlow, a deep learning framework that uses guided randomized search of the SOAP space to find a fast parallelization strategy for a specific parallel machine. To accelerate this search, FlexFlow introduces a novel execution simulator that can accurately predict a parallelization strategy's performance and is three orders of magnitude faster than prior approaches that have to execute each strategy. We evaluate FlexFlow with six real-world DNN benchmarks on two GPU clusters and show that FlexFlow can increase training throughput by up to 3.8x over state-of-the-art approaches, even when including its search time, and also improves scalability.