Supporting Very Large Models using Automatic Dataflow Graph Partitioning

By Minjie Wang, Chien-chin Huang, Jinyang Li. 2018

175 citations recorded by the source corpus.

This paper presents Tofu, a system that partitions very large DNN models across multiple GPU devices to reduce per-GPU memory footprint. Tofu is designed to partition a dataflow graph of fine-grained tensor operators used by platforms like MXNet and TensorFlow. In order to automatically partition each operator, we propose to describe the semantics of an operator in a simple language inspired by Halide. To optimally partition different operators in a dataflow graph, Tofu uses a recursive search algorithm that minimizes the total communication cost. Our experiments on an 8-GPU machine show that Tofu enables the training of very large CNN and RNN models. It also achieves 25% - 400% speedup over alternative approaches to train very large models.

Subfields: Computer Vision and Pattern Recognition, Hardware and Architecture, Model Architectures, Efficiency & Compression

How researchers see this paper

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

In Computer Vision and Pattern Recognition

In Hardware and Architecture

In Model Architectures

In Efficiency & Compression

Citations in the corpus

ZeRO: Memory optimizations Toward Training Trillion Parameter Models (2019, background): “When a model does not fit in the device memory, model parallelism (MP) [7], [4], [11] and pipeline parallelism (PP) [12], [13] split the model among processes, in vertical and horizontal way respectively.”
GSPMD: General and Scalable Parallelization for ML Computation Graphs (2021, background): “Tofu [39] is another parallelization system for large models, but it supports only limited partition strategies (e.”
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).”
ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep learning (2021, background): “When models do not fit in device memory, model parallelism3 (MP) [7, 17, 18] and”
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.”
Wireless Network Intelligence at the Edge (2018, neutral): “In order to minimize the dependence among the split segments, one can first unroll the original NN model and construct its dataflow graph [107], [108].”
Pre-Trained Models: Past, Present and Future (2021, methodology): “, 2019), which were designed in the early days without initially foreseeing the emerging requirements such as model/pipeline parallelism of large models (Brown et al., 2020; Huang et al., 2019b; Wang et al., 2019).”
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”
DAPPLE: a pipelined data parallel approach for training large models (2020, background): “Some researchers have been seeking for the optimal placement strategy to assign operations in a DNN to different devices [57]–[59] to further improve system efficiency.”
SwapAdvisor: Pushing Deep Learning Beyond the GPU Memory Limit via Smart Swapping (2020, background): “[22, 23, 50] propose to auto-mate the model parallelism and reduce the communication with dataflow graph analysis.”
AccPar: Tensor Partitioning for Heterogeneous Deep Learning Accelerators (2020, methodology): “Recursive methods [74,118] are proposed for tensor partitioning on multiple devices and dynamic programming methods [116,118] are proposed for tensor partitioning layer-wisely.”
HetPipe: Enabling Large DNN Training on (Whimpy) Heterogeneous GPU Clusters through Integration of Pipelined Model Parallelism and Data Parallelism (2020, background): “, the number of parameters) have continuously been increasing in order to improve the accuracy and quality of models and to deal with complex features of data [17, 42, 48, 49].”
Buddy Compression: Enabling Larger Memory for Deep Learning and HPC Workloads on GPUs (2019, background): “Today, applications with large memory footprints must either (i) scale out to many GPUs for capacity purposes (inefficient resource utilization) [5, 6], (ii) explicitly orchestrate NVLink2”
A generic communication scheduler for distributed DNN training acceleration (2019, background): “, by grouping or coalescing multiple operators or tensors [36]).”
This paper is included in the Proceedings of the 16th USENIX Symposium on Operating Systems Design and Implementation. (unknown): “For compute costs c v , we set all of them as zero following the same motivation in [55].”
Gradient Compression Supercharged High-Performance Data Parallel DNN Training (2021, methodology): “Model Parallelism [15, 65] and Pipeline Parallelism [48] are often combined with Data Parallelism for large-scale deployment [71, 72], which can benefit from HiPress .”