Source-linked AI summary
Beyond Data and Model Parallelism for Deep Neural Networks
Zhihao Jia, Matei Zaharia, Alex Aiken
TL;DR
Training DNNs increasingly requires distributed parallelism, while existing data/model parallelism can be suboptimal. FlexFlow searches the broader SOAP space using guided randomized search and an execution simulator, outperforming existing approaches on six DNN benchmarks across two GPU clusters.
Problem
Growing DNN models and datasets make distributed training standard, but existing systems mainly support batch-dimension data parallelism and lack flexible operation-level parallelization.
Method
FlexFlow defines the SOAP search space and uses guided randomized search with an accurate, three-orders-of-magnitude-faster execution simulator to find efficient strategies.
Results
Across six real-world DNN benchmarks on two GPU clusters, FlexFlow outperforms data/model parallelism and expert-designed strategies by up to 3.3× in training throughput and achieves significantly better scaling.
Takeaways & Limitations
FlexFlow provides an automated way to find efficient DNN parallelization strategies over a broader space than prior systems.
Takeaways & Limitations
The approach may not apply to applications whose operation execution times depend on input data, although the studied dense-matrix DNNs have predictable execution times.
Abstract
from arXiv · showhide
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.
1 Introduction
DNN training increasingly requires distributed parallelism, but common data/model strategies and limited automated methods can miss efficient configurations. FlexFlow addresses this with the broader SOAP space, a fast execution simulator, and guided search.
- Growing models and datasets have made distributed training across heterogeneous clusters standard practice.
- Data parallelism replicates the network on each device, while model parallelism partitions operations across devices and restricts within-operation parallelism.Data parallelism synchronizes parameters; model parallelism instead requires transfers between operations.
- Expert-designed and automated approaches improve on basic strategies but remain limited in their supported parallelization spaces.Prior automated systems use restricted searches, while expert strategies can still be suboptimal.
- SOAP covers Sample, Operation, Attribute, and Parameter parallelism, including combinations that existing systems may omit.FlexFlow therefore searches a more comprehensive space containing existing approaches as special cases.
- Three orders of magnitude faster than profiling real executions, FlexFlow’s simulator predicts strategy performance using operator measurements and incremental delta simulation.It measures each operator type and input size, then reuses those estimates across strategies.
- Across six DNN benchmarks, simulated execution times differed from measured times by less than 30% and preserved the ordering of candidate strategies.
- FlexFlow uses MCMC-guided search over SOAP candidates and reports throughput gains of up to 3.3× over data/model and expert-designed strategies.It also outperformed REINFORCE by 3.4–3.8× and OptCNN by 1.2–1.6× on the reported hardware configurations.
2 Related Work
Prior systems mainly use data or model parallelism, while expert-designed and automated methods explore narrower alternatives. FlexFlow broadens coverage across SOAP dimensions, including hybrid within-operation parallelism.
- Data parallelism replicates an entire DNN and can bottleneck scalability for operations with many parameters.
- Model parallelism assigns disjoint DNN subsets to devices, reducing parameter synchronization but exposing limited parallelism.
- Expert-designed strategies switch between data and model parallelism for different layers but remain tailored to specific DNN patterns.
- REINFORCE searches device placement for model parallelism, whereas OptCNN targets within-operation parallelism for linear computation graphs.
- Existing approaches cover subsets of SOAP dimensions, while FlexFlow also supports hybrid combinations of sample, attribute, and parameter dimensions.
3 Overview
FlexFlow maps an operator graph onto a hardware topology, searches parallelization strategies with simulated execution, and returns a strategy suited to that configuration. Its main scope boundary is the simulator’s assumption of predictable, input-independent operation times.
- Programming Interface: FlexFlow represents DNN computation and state as an operator graph whose nodes are operations and edges are tensors.
- Programming Interface: Unlike systems that default to data parallelism or require manual placement, FlexFlow automatically searches for a strategy for the operator graph and device topology.
- Programming Interface: The device topology records available CPU/GPU devices and their interconnections, including connection bandwidth and latency.
- Overview: The framework aims to improve programmability for complex graphs and portability by selecting strategies for each hardware configuration without application changes.
- Overview: FlexFlow’s optimizer uses MCMC proposals and delta-simulated execution estimates to generate an efficient parallelization strategy.
- Limitations: The simulator may not apply to applications whose operation times depend on input contents, although the studied dense-matrix DNNs satisfy its predictability assumption.
- Parallelizable Dimensions: Table 1 identifies parallelizable dimensions for operations, including sample, channel, length, and image-position dimensions.
- Parallelization Configurations: Figure 3 illustrates single-dimension and multidimensional partitioning configurations for 1D convolution using dashed tensor-partition boundaries.
4 The SOAP Search Space
SOAP defines DNN parallelization across Sample, Operation, Attribute, and Parameter dimensions by partitioning operation outputs and assigning resulting tasks to devices.
- SOAP models operation parallelization by partitioning each operation’s output tensor across divisible dimensions.
- The sample dimension is always parallelizable; other dimensions are attributes unless partitioning them requires splitting model parameters.
- A configuration specifies positive parallelism degrees for selected dimensions, whose product determines the number of independent tasks.
- Equal-size partitions balance workloads, while each task receives a device assignment and derives required inputs from its output sub-tensor.
- A parallelization strategy independently selects one configuration for every operation, including combinations of dimensions such as those illustrated for matrix multiplication.
5 Execution Simulator
FlexFlow’s execution simulator predicts strategy runtime by constructing task graphs and scheduling computation and communication tasks under explicit execution assumptions. Delta simulation reuses prior timelines to accelerate evaluation of nearby strategies.
- The simulator predicts execution time from an operator graph, device topology, and parallelization strategy instead of measuring each real execution.It runs up to three orders of magnitude faster than real executions.
- The simulator assumes predictable task times independent of input contents, bandwidth-limited transfers, FIFO device scheduling, and negligible runtime overhead.
- Task Graph: Task-graph construction creates normal tasks from operation configurations and dependency edges, adding communication tasks when shared tensors cross devices.Edges encode ordering constraints, while communication tasks represent data transfers.
- Task execution times are measured for normal tasks and cached by operation type and output size; communication times are estimated from tensor size and bandwidth.
- Delta Simulation Algorithm: Delta simulation updates only tasks affected by a changed operation configuration, propagates timing changes through the task graph, and produces the same timeline as full simulation.Because MCMC changes one operation at a time, most of the previous timeline commonly remains unchanged.
6 Execution Optimizer
FlexFlow’s optimizer treats parallelization as a cost-minimization problem and heuristically searches the exponentially large SOAP space with MCMC, accepting some worse candidates to escape local minima.
- The optimizer uses the simulator as an oracle to minimize predicted execution time over parallelization strategies.
- Optimal strategy search is NP-hard, and the number of possible strategies grows exponentially with the operator graph’s operations.
- MCMC Sampling: MCMC converts strategy cost into a probability distribution and samples strategies so higher-probability, lower-cost candidates are visited more often.
- MCMC Sampling: Metropolis-Hastings maintains a current strategy, proposes a modified strategy, and accepts it according to the proposal and cost-based criterion.
- MCMC Sampling: Higher-cost proposals may still be accepted with decreasing probability, allowing the search to escape local minima while preferring lower-cost strategies.
- FlexFlow changes one operation’s configuration per proposal and searches from existing or random initial strategies until a time or improvement-stagnation criterion is met.
7 FlexFlow Runtime
FlexFlow’s runtime supports parallelizing individual operations across arbitrary combinations of parallelizable dimensions, extending beyond the limited capabilities of existing systems.
- Existing deep learning systems primarily support operation parallelization in the batch dimension through data parallelism.Parallelizing other dimensions or combinations is non-trivial, and operation-level control is generally unavailable.
- FlexFlow implements distributed execution using Legion, with cuDNN and cuBLAS for DNN operations and high-dimensional partitioning for supported dimensions.
- The runtime can parallelize an operation in any combination of parallelizable dimensions and control parallelization at individual-operation granularity.
8 Evaluation
FlexFlow is evaluated across six DNN benchmarks and two GPU clusters against data parallelism, expert-designed strategies, REINFORCE, and OptCNN. It generally improves parallelization performance while preserving model accuracy and reducing communication or computation costs.
- 8.2.1 Per-iteration Performance: 1.3-3.3× speedup over baselines is achieved on all benchmarks except ResNet-101, where FlexFlow finds a strategy similar to data parallelism.The improvement comes from faster parallelization strategies rather than different DNN computations.
- 8.2.1 Per-iteration Performance: 2-5.5× lower per-iteration data transfers are achieved for NMT on 64 K80 GPUs compared with other parallelization approaches.FlexFlow supports overlapping transfers with computation, but reducing communication becomes increasingly beneficial as device count grows.
- 8.2.1 Per-iteration Performance: 20% lower overall task computation time than data parallelism is achieved for NMT, while retaining intra-operation parallelism and load balance.An expert-designed strategy has slightly lower task computation time but worse execution performance because model parallelism disables intra-operation parallelism and causes imbalance.
- 8.2.1 Per-iteration Performance: State-of-the-art benchmark accuracies are maintained because FlexFlow performs the same computation as other deep learning systems.The accuracy claim is verified on the DNN benchmarks used in the experiments.
- 8.2.3 Automated Parallelization Optimizer: 3.4-3.8× speedup over REINFORCE is achieved by FlexFlow strategies, which the authors attribute to exploring a larger search space.FlexFlow’s optimizer finds strategies in 14-40 seconds on one compute node, whereas REINFORCE takes 12-27 hours and uses up to 160 compute nodes.
- 8.2.2 Comparison with OptCNN: 1.2-1.6× speedup over OptCNN is achieved for DNNs with non-linear operator graphs by exploiting parallelism across different operations.The frameworks find the same strategies for AlexNet and ResNet, but different strategies for the other evaluated DNNs.
8.3 Execution Simulator
The execution simulator is evaluated for prediction accuracy and search efficiency. It preserves the ordering of actual strategy performance while delta simulation substantially reduces optimizer time.
- 8.3 Execution Simulator: The simulator evaluation uses simulator accuracy and simulator execution time as its two metrics.
- 8.3.1 Simulator Accuracy: Simulated execution times preserve the actual ordering of strategies sharing an operator graph and device topology.Predicted and measured times fall within the figure’s stated 0% and 30% relative-difference bounds.
- 8.3.2 Simulator Efficiency: 16 minutes versus 6 minutes is the termination time for full versus delta simulation when searching for an NMT strategy on 16 P100 GPUs.With a time budget below 8 minutes, delta simulation finds a better strategy than full simulation.
- 8.3.2 Simulator Efficiency: 2.2-6.9× faster optimizer execution is achieved by delta simulation compared with full simulation.The speedup increases as the number of devices scales.
8.4 Search Algorithm
FlexFlow’s search algorithm is tested against global and local optimality. It finds global optima in the small evaluated spaces and locally optimal strategies in larger spaces.
- Global Optimality: FlexFlow finds the global optimal strategy for both LeNet and a restricted RNNLM search space of approximately 10^11 strategies.The exhaustive optimality searches use four devices and take 0.8 hours for LeNet and 18 hours for RNNLM.
- Local Optimality: All strategies returned by FlexFlow are locally optimal across six DNNs evaluated on 2, 4, and 8 devices.Local optimality is tested by exhaustively comparing each discovered strategy with all of its neighbors.
8.5 Case Studies
The case studies show FlexFlow adapting parallelization to operation structure and device topology. Its discovered strategies balance workload and reduce communication costs relative to data parallelism.
- Inception-v3: 12% lower per-iteration execution time and 75% lower parameter-synchronization cost are achieved for Inception-v3 on four P100 GPUs.The strategy combines intra-operation parallelism on critical-path operations with intra- and inter-operation parallelism across branches.
- Inception-v3: Adjacent GPUs with direct connections are favored for Inception-v3 on four asymmetric K80 connections to reduce communication costs.
- NMT: Different NMT layers use different parallelization strategies rather than one uniform strategy across the model.The passage introduces layer-specific choices, including using fewer GPUs for parameter-heavy, computation-light layers to reduce synchronization costs.
9 Conclusion
FlexFlow automatically searches for efficient DNN parallelization strategies using guided randomized search and an execution simulator, outperforming state-of-the-art approaches on six benchmarks across two GPU clusters.
- FlexFlow uses guided randomized search and an execution simulator to find efficient parallelization strategies for DNN applications.The simulator predicts DNN performance efficiently and accurately during the search.
- FlexFlow significantly outperforms state-of-the-art parallelization approaches across six real-world DNN benchmarks on two GPU clusters.