# Automatic Graph Partitioning for Very Large-scale Deep Learning

arXiv: https://arxiv.org/abs/2103.16063 · HTML: https://mlsanity.com/paper/2103.16063/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 29 references that are in the map.
- First posted on arXiv 30 Mar 2021.

## Builds on (most cited first, 21 of 21)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [Big Transfer (BiT): General Visual Representation Learning](https://mlsanity.com/paper/1912.11370/) (2019)
- [Beyond Data and Model Parallelism for Deep Neural Networks](https://mlsanity.com/paper/1807.05358/) (2018)
- [Mesh-TensorFlow: Deep Learning for Supercomputers](https://mlsanity.com/paper/1811.02084/) (2018)
- [Memory-Efficient Pipeline-Parallel DNN Training](https://mlsanity.com/paper/2006.09503/) (2020)
- [PipeDream: Fast and Efficient Pipeline Parallel DNN Training](https://mlsanity.com/paper/1806.03377/) (2018)
- [HetPipe: Enabling Large DNN Training on (Whimpy) Heterogeneous GPU Clusters through Integration of Pipelined Model Parallelism and Data Parallelism](https://mlsanity.com/paper/2005.14038/) (2020)
- [Supporting Very Large Models using Automatic Dataflow Graph Partitioning](https://mlsanity.com/paper/1807.08887/) (2018)
- [Efficient and Robust Parallel DNN Training through Model Parallelism on Multi-GPU Platform](https://mlsanity.com/paper/1809.02839/) (2018)
- [torchgpipe: On-the-fly Pipeline Parallelism for Training Giant Models](https://mlsanity.com/paper/2004.09910/) (2020)
- [Exploring Hidden Dimensions in Parallelizing Convolutional Neural Networks](https://mlsanity.com/paper/1802.04924/) (2018)
- [XPipe: Efficient Pipeline Model Parallelism for Multi-GPU DNN Training](https://mlsanity.com/paper/1911.04610/) (2019)
- [AMPNet: Asynchronous Model-Parallel Training for Dynamic Neural Networks](https://mlsanity.com/paper/1705.09786/) (2017)

## Built on by (most cited first, 7 of 7)

- [A Hybrid Tensor-Expert-Data Parallelism Approach to Optimize Mixture-of-Experts Training](https://mlsanity.com/paper/2303.06318/) (2023)
- [AxoNN: An asynchronous, message-driven parallel framework for extreme-scale deep learning](https://mlsanity.com/paper/2110.13005/) (2021)
- [FlexPipe: Adapting Dynamic LLM Serving Through Inflight Pipeline Refactoring in Fragmented Serverless Clusters](https://mlsanity.com/paper/2510.11938/) (2025)
- [Model Parallelism on Distributed Infrastructure: A Literature Review from Theory to LLM Case-Studies](https://mlsanity.com/paper/2403.03699/) (2024)
- [DISCO: Distributed Inference with Sparse Communications](https://mlsanity.com/paper/2302.11180/) (2023)
- [Balanced and Elastic End-to-end Training of Dynamic LLMs](https://mlsanity.com/paper/2505.14864/) (2025)
- [Probabilistic Algorithms for Ising Machines from Optimization to Generative AI](https://mlsanity.com/paper/2610.03972/) (2026)

## Abstract (arXiv)

This work proposes RaNNC (Rapid Neural Network Connector) as middleware for automatic hybrid parallelism. In recent deep learning research, as exemplified by T5 and GPT-3, the size of neural network models continues to grow. Since such models do not fit into the memory of accelerator devices, they need to be partitioned by model parallelism techniques. Moreover, to accelerate training for huge training data, we need a combination of model and data parallelisms, i.e., hybrid parallelism. Given a model description for PyTorch without any specification for model parallelism, RaNNC automatically partitions the model into a set of subcomponents so that (1) each subcomponent fits a device memory and (2) a high training throughput for pipeline parallelism is achieved by balancing the computation times of the subcomponents. In our experiments, we compared RaNNC with two popular frameworks, Megatron-LM (hybrid parallelism) and GPipe (originally proposed for model parallelism, but a version allowing hybrid parallelism also exists), for training models with increasingly greater numbers of parameters. In the pre-training of enlarged BERT models, RaNNC successfully trained models five times larger than those Megatron-LM could, and RaNNC's training throughputs were comparable to Megatron-LM's when pre-training the same models. RaNNC also achieved better training throughputs than GPipe on both the enlarged BERT model pre-training (GPipe with hybrid parallelism) and the enlarged ResNet models (GPipe with model parallelism) in all of the settings we tried. These results are remarkable, since RaNNC automatically partitions models without any modification to their descriptions; Megatron-LM and GPipe require users to manually rewrite the models' descriptions.