Beyond Data and Model Parallelism for Deep Neural Networks

By Matei A. Zaharia, A. Aiken. 2018

624 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Hardware and Architecture, Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Hardware and Architecture

In Model Architectures

Citations in the corpus

Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (2019, methodology): “FlexFlow (Jia et al., 2018), a deep learning framework orchestrating”
Mesh-TensorFlow: Deep Learning for Supercomputers (2018, methodology): “[20, 19] implemented a framework that uses cost modeling to pick the best parallelization strategy, including how to partition work for each operation.”
GSPMD: General and Scalable Parallelization for ML Computation Graphs (2021, methodology): “FlexFlow [20] uses automated search to discover the partitioning strategy of operators in a graph for better performance.”
PanGu-α: Large-scale Autoregressive Pretrained Chinese Language Models with Auto-parallel Computation (2021, methodology): “One regime is op-level model parallelism [17, 18, 19, 20, 21, 22, 23], which partitions its involved tensors of each operator (layer), as shown in Figure 4(b).”
Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (2021, methodology): “FlexFlow [17], PipeDream [22], Tarnawski et al.”
GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding (2020, methodology): “For example, pipelining [15, 24] splits a large model’s layers into multiple stages, while operator-level partitioning [23, 81] splits individual operators into smaller parallel operators.”
Memory-Efficient Pipeline-Parallel DNN Training (2020, methodology): “PipeDream-2BW ’s planner considers the impact of these decisions on both throughput and memory footprint, unlike PipeDream and FlexFlow.”
Reducing Activation Recomputation in Large Transformer Models (2022, methodology): “In other words, our technique, which mixes both tensor and sequence parallelism, reduces the activation memory signiﬁcantly without any additional compute, communication, or memory overhead.”
PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel (2023, methodology): “Further, compiler-based techniques such as Alpa [35], GSPMD [31], and FlexFlow [9] leverage profiling, performance modeling, user annotations and search to find the best configuration across the parallelism space of data, tensor and pipeline for a given cluster.”
Pre-Trained Models: Past, Present and Future (2021, methodology): “, 2018), FlexFlow (Jia et al., 2019), OneFlow (one, 2021), MindSpore (min, 2021) and GShard (Lepikhin et al.”
Supporting Very Large Models using Automatic Dataflow Graph Partitioning (2018, methodology): “As such, they either develop specialized implementation for specific models [9, 13] or allow only a composition of common DNN layers [8, 10, 14, 15].”
Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization (2019, neutral): “However, model parallelism requires access to additional expensive compute accelerators, fast networks, and non-trivial partitioning of model state to balance communication and computation [18, 31, 39].”
TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models (2021, background): “Many existing works (Jia et al., 2018; 2019; Wang et al., 2019; Shazeer et al., 2018) study how to optimize the partitioning schemes for different operations to maximize throughput and minimize communication overheads, among which, Megatron-LM (Figure 1b; Shoeybi et al., 2019) designs partitioning…”
Automatic Cross-Replica Sharding of Weight Update in Data-Parallel Training (2020, methodology): “FlexFlow [12] uses automated search to discover the optimal partition of operations in a graph, which has a flexible search space to cover model and data parallelism.”
DAPPLE: a pipelined data parallel approach for training large models (2020, methodology): “Model Parallelism [50] partitions DNN models among GPUs to mitigate communication overhead and memory bottlenecks for distributed training [10], [14], [39], [40], [51]–[54].”
SwapAdvisor: Pushing Deep Learning Beyond the GPU Memory Limit via Smart Swapping (2020, methodology): “[22, 23, 50] propose to automate the model parallelism and reduce the communication with dataflow graph analysis.”
Understanding Capacity-Driven Scale-Out Neural Recommendation Inference (2020, background): “…for model-parallelism in TensorFlow–and in turn huge models–but have focused on language models and training which don’t have the variety and number of features shown in Figure 5 and emphasize training throughput compared to inference’s latency-bounded throughput [15], [36], [37], [40], [41].”
Efficient and Robust Parallel DNN Training through Model Parallelism on Multi-GPU Platform (2018, background): “(Jia et al., 2019; 2018) partition the tensor along multiple dimensions and search for best parallelization strategy.”