[Paper Review] Beyond Data and Model Parallelism for Deep Neural Networks

Zhihao Jia, Matei Zaharia+1 more|arXiv (Cornell University)|Jul 14, 2018

Advanced Neural Network Applications31 references147 citations

TL;DR

FlexFlow defines a broader SOAP space (Sample-Operation-Attribute-Parameter) for DNN parallelism, uses a fast execution simulator with MCMC search to find efficient strategies, and achieves substantial throughput gains over prior methods.

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.

Motivation & Objective

• Motivate the need for more comprehensive parallelization beyond data and model parallelism.
• Formalize a broader SOAP search space that includes Sample, Operation, Attribute, and Parameter dimensions.
• Develop a fast execution simulator to predict performance and guide optimization.
• Propose FlexFlow, a framework that automatically discovers and executes fast parallelization strategies.
• Demonstrate throughput and scalability improvements on real-world DNN benchmarks.

Proposed method

• Define the SOAP search space (Sample, Operation, Attribute, Parameter) for parallelizing DNNs across devices.
• Develop a fast execution simulator that predicts performance with low variance and high speed, enabling broad search.
• Use a Markov Chain Monte Carlo (MCMC) optimizer to explore SOAP strategies based on simulated performance.
• Implement a full and delta simulation algorithm to efficiently evaluate strategy changes.
• Build a distributed runtime (Legion) to execute the discovered parallelization strategies.

Experimental results

Research questions

• RQ1 Can the SOAP space yield faster parallelizations than traditional data/model parallelism and expert-designed strategies?
• RQ2 How accurate and fast is the FlexFlow execution simulator compared with real executions?
• RQ3 What throughput and scalability gains are achievable on real-world DNN benchmarks across GPU clusters?
• RQ4 How does FlexFlow compare to REINFORCE and OptCNN in discovering efficient strategies?
• RQ5 What is the impact of broader parallelization on communication costs and scheduling?

Key findings

• FlexFlow increases training throughput by up to 3.8x over state-of-the-art approaches.
• The simulator-guided search achieves 14–40 seconds to find strategies that REINFORCE needs 12–27 hours for on 4 GPUs and 160 nodes total.
• FlexFlow achieves up to 3.3x higher throughput and up to 5x reduction in communication costs in evaluations.
• FlexFlow outperforms REINFORCE by 3.4–3.8x on the same hardware configuration and outperforms OptCNN by 1.2–1.6x by supporting a broader SOAP space.
• The simulator's relative error vs real execution time is below 30% across measured executions, preserving time ordering of strategies.