# Baechi: Fast Device Placement of Machine Learning Graphs

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 77 references that are in the map.
- First posted on arXiv 20 Jan 2023.

## Builds on (most cited first, 30 of 33)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [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)
- [Rethinking the Inception Architecture for Computer Vision](https://mlsanity.com/paper/1512.00567/) (2015)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Google's Neural Machine Translation System: Bridging the Gap between Human and Machine Translation](https://mlsanity.com/paper/1609.08144/) (2016)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [Towards Federated Learning at Scale: System Design](https://mlsanity.com/paper/1902.01046/) (2019)
- [Caffe: Convolutional Architecture for Fast Feature Embedding](https://mlsanity.com/paper/1408.5093/) (2014)
- [One weird trick for parallelizing convolutional neural networks](https://mlsanity.com/paper/1404.5997/) (2014)
- [MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems](https://mlsanity.com/paper/1512.01274/) (2015)
- [Beyond Data and Model Parallelism for Deep Neural Networks](https://mlsanity.com/paper/1807.05358/) (2018)
- [PipeDream: Fast and Efficient Pipeline Parallel DNN Training](https://mlsanity.com/paper/1806.03377/) (2018)
- [DAPPLE: A Pipelined Data Parallel Approach for Training Large Models](https://mlsanity.com/paper/2007.01045/) (2020)
- [Device Placement Optimization with Reinforcement Learning](https://mlsanity.com/paper/1706.04972/) (2017)
- [TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models](https://mlsanity.com/paper/2102.07988/) (2021)
- [Theano: A Python framework for fast computation of mathematical expressions](https://mlsanity.com/paper/1605.02688/) (2016)
- [PipeMare: Asynchronous Pipeline Parallel DNN Training](https://mlsanity.com/paper/1910.05124/) (2019)
- [Supporting Very Large Models using Automatic Dataflow Graph Partitioning](https://mlsanity.com/paper/1807.08887/) (2018)
- [Petuum: A New Platform for Distributed Machine Learning on Big Data](https://mlsanity.com/paper/1312.7651/) (2013)
- [Machine learning for Internet of Things data analysis: A survey](https://mlsanity.com/paper/1802.06305/) (2018)
- [Ordering Chaos: Memory-Aware Scheduling of Irregularly Wired Neural Networks for Edge Devices](https://mlsanity.com/paper/2003.02369/) (2020)
- [Placeto: Learning Generalizable Device Placement Algorithms for Distributed Machine Learning](https://mlsanity.com/paper/1906.08879/) (2019)
- [Efficient Algorithms for Device Placement of DNN Graph Operators](https://mlsanity.com/paper/2006.16423/) (2020)
- [Nimble: Lightweight and Parallel GPU Task Scheduling for Deep Learning](https://mlsanity.com/paper/2012.02732/) (2020)
- [Characterizing Deep Learning Training Workloads on Alibaba-PAI](https://mlsanity.com/paper/1910.05930/) (2019)
- [Amazon SageMaker Model Parallelism: A General and Flexible Framework for Large Model Training](https://mlsanity.com/paper/2111.05972/) (2021)

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

- [Moirai: Towards Optimal Placement for Distributed Inference on Heterogeneous Devices](https://mlsanity.com/paper/2312.04025/) (2023)
- [Expediting Distributed DNN Training with Device Topology-Aware Graph Deployment](https://mlsanity.com/paper/2302.06126/) (2023)

## Abstract (arXiv)

Machine Learning graphs (or models) can be challenging or impossible to train when either devices have limited memory, or models are large. To split the model across devices, learning-based approaches are still popular. While these result in model placements that train fast on data (i.e., low step times), learning-based model-parallelism is time-consuming, taking many hours or days to create a placement plan of operators on devices. We present the Baechi system, the first to adopt an algorithmic approach to the placement problem for running machine learning training graphs on small clusters of memory-constrained devices. We integrate our implementation of Baechi into two popular open-source learning frameworks: TensorFlow and PyTorch. Our experimental results using GPUs show that: (i) Baechi generates placement plans 654 X - 206K X faster than state-of-the-art learning-based approaches, and (ii) Baechi-placed model's step (training) time is comparable to expert placements in PyTorch, and only up to 6.2% worse than expert placements in TensorFlow. We prove mathematically that our two algorithms are within a constant factor of the optimal. Our work shows that compared to learning-based approaches, algorithmic approaches can face different challenges for adaptation to Machine learning systems, but also they offer proven bounds, and significant performance benefits.