Beyond Data and Model Parallelism for Deep Neural Networks

TL;DR

This work introduces FlexFlow, a framework that generalizes DNN parallelism with the SOAP (Sample-Operation-Attribute-Parameter) space and auto-discovers fast strategies on heterogeneous clusters. A fast delta-based execution simulator serves as an oracle for a guided MCMC optimizer, enabling scalable search across a much larger design space than prior systems. Empirical evaluations on six real-world benchmarks across two GPU clusters show substantial throughput improvements (up to 3.8×) and reduced communication, outperforming data/model parallelism, expert-designed baselines, and existing automated frameworks. The approach enables more flexible, scalable, and hardware-tuned parallelization for diverse DNN workloads, with a practical runtime that adapts strategies to specific cluster topologies.

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.

Figures (14)

• Figure 1: The parallelism dimensions explored by different approaches. S, O, A, and P indicate parallelism in the Sample, Operation, Attribute, and Parameter dimensions (see Section \ref{['sec:problem']}). Hybrid parallelism shows if an approach supports parallelizing an operation in a combination of the sample, attribute, and parameter dimensions (see FIgure \ref{['fig:example_parallelization']}). OptCNN is designed for DNNs with linear computation graphs.
• Figure 2: FlexFlow overview.
• Figure 3: Example parallelization configurations for 1D convolution. Dashed lines show partitioning the tensor.
• Figure 4: An example parallelization configuration for a matrix multiplication operation.
• Figure 5: Simulating an example parallelization strategy. The tasks' exeTime and device are shown on the top of each column. In Figure \ref{['fig:example_simulation_base']} and \ref{['fig:example_simulation_opt']}, the word "r" and "s" indicate the readyTime and startTime of each task, respectively, and the dashed edges represents the nextTask.