Supporting Very Large Models using Automatic Dataflow Graph Partitioning

Abstract

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 in order to work transparently with a general-purpose deep learning platform like MXNet. In order to automatically partition each operator, we propose to describe the semantics of an operator in a simple language which represents tensors as lambda functions mapping from tensor coordinates to values. 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.

arXiv commentRevision for Eurosys'19

References in corpus (19)

• Adam: A Method for Stochastic Optimization
• Sequence to Sequence Learning with Neural Networks
• Google's Neural Machine Translation System: Bridging the Gap between Human and Machine Translation
• MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems
• Proceedings of the 29th International Conference on Machine Learning (ICML-12)
• Compressing Deep Convolutional Networks using Vector Quantization
• One weird trick for parallelizing convolutional neural networks
• Binarized Neural Networks
• Exploring the Limits of Language Modeling
• Learning both Weights and Connections for Efficient Neural Networks
• Training Deep Nets with Sublinear Memory Cost
• Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions
• Device Placement Optimization with Reinforcement Learning
• Supporting Very Large Models using Automatic Dataflow Graph Partitioning
• Beyond Data and Model Parallelism for Deep Neural Networks
• Memory-Efficient Backpropagation Through Time
• Exploring Hidden Dimensions in Parallelizing Convolutional Neural Networks
• A Systematic Approach to Blocking Convolutional Neural Networks
• Profile-guided memory optimization for deep neural networks

Cited by in corpus (29)

• Supporting Very Large Models using Automatic Dataflow Graph Partitioning
• PanGu-$\alpha$: Large-scale Autoregressive Pretrained Chinese Language Models with Auto-parallel Computation
• ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning
• Galvatron: Efficient Transformer Training over Multiple GPUs Using Automatic Parallelism
• HybridFlow: A Flexible and Efficient RLHF Framework
• TensorOpt: Exploring the Tradeoffs in Distributed DNN Training with Auto-Parallelism
• OneFlow: Redesign the Distributed Deep Learning Framework from Scratch
• TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models
• DAPPLE: A Pipelined Data Parallel Approach for Training Large Models
• HongTu: Scalable Full-Graph GNN Training on Multiple GPUs (via communication-optimized CPU data offloading)
• ReCycle: Resilient Training of Large DNNs using Pipeline Adaptation
• HAP: SPMD DNN Training on Heterogeneous GPU Clusters with Automated Program Synthesis
• Pre-Trained Models: Past, Present and Future
• DistTrain: Addressing Model and Data Heterogeneity with Disaggregated Training for Multimodal Large Language Models
• Improving Automatic Parallel Training via Balanced Memory Workload Optimization
• Orchestrating the Development Lifecycle of Machine Learning-Based IoT Applications: A Taxonomy and Survey
• HyPar: Towards Hybrid Parallelism for Deep Learning Accelerator Array
• AutoDDL: Automatic Distributed Deep Learning with Near-Optimal Bandwidth Cost
• PaSE: Parallelization Strategies for Efficient DNN Training
• ActNN: Reducing Training Memory Footprint via 2-Bit Activation Compressed Training
• A Computational-Graph Partitioning Method for Training Memory-Constrained DNNs
• Heterogeneity-Aware Asynchronous Decentralized Training
• Mist: Efficient Distributed Training of Large Language Models via Memory-Parallelism Co-Optimization
• Parallax: Sparsity-aware Data Parallel Training of Deep Neural Networks
• Scaling Deep Learning Computation over the Inter-Core Connected Intelligence Processor with T10
• Automatic Graph Partitioning for Very Large-scale Deep Learning
• Scaling Distributed Deep Learning Workloads beyond the Memory Capacity with KARMA
• Training on the Edge: The why and the how
• TAPAS: Fast and Automatic Derivation of Tensor Parallel Strategies for Large Neural Networks