# DiviML: A Module-based Heuristic for Mapping Neural Networks onto Heterogeneous Platforms

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 17 of its 42 references that are in the map.
- First posted on arXiv 31 Jul 2023.

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

- [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)
- [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)
- [Going Deeper with Convolutions](https://mlsanity.com/paper/1409.4842/) (2014)
- [Neural Architecture Search with Reinforcement Learning](https://mlsanity.com/paper/1611.01578/) (2016)
- [MnasNet: Platform-Aware Neural Architecture Search for Mobile](https://mlsanity.com/paper/1807.11626/) (2018)
- [ProxylessNAS: Direct Neural Architecture Search on Target Task and Hardware](https://mlsanity.com/paper/1812.00332/) (2018)
- [In-Datacenter Performance Analysis of a Tensor Processing Unit](https://mlsanity.com/paper/1704.04760/) (2017)
- [BRP-NAS: Prediction-based NAS using GCNs](https://mlsanity.com/paper/2007.08668/) (2020)
- [Exploring Randomly Wired Neural Networks for Image Recognition](https://mlsanity.com/paper/1904.01569/) (2019)
- [Device Placement Optimization with Reinforcement Learning](https://mlsanity.com/paper/1706.04972/) (2017)
- [Torch.fx: Practical Program Capture and Transformation for Deep Learning in Python](https://mlsanity.com/paper/2112.08429/) (2021)
- [Ordering Chaos: Memory-Aware Scheduling of Irregularly Wired Neural Networks for Edge Devices](https://mlsanity.com/paper/2003.02369/) (2020)
- [DLA: Compiler and FPGA Overlay for Neural Network Inference Acceleration](https://mlsanity.com/paper/1807.06434/) (2018)
- [A Computational-Graph Partitioning Method for Training Memory-Constrained DNNs](https://mlsanity.com/paper/2008.08636/) (2020)
- [DYNAMAP: Dynamic Algorithm Mapping Framework for Low Latency CNN Inference](https://mlsanity.com/paper/2012.00912/) (2020)

## Abstract (arXiv)

Datacenters are increasingly becoming heterogeneous, and are starting to include specialized hardware for networking, video processing, and especially deep learning. To leverage the heterogeneous compute capability of modern datacenters, we develop an approach for compiler-level partitioning of deep neural networks (DNNs) onto multiple interconnected hardware devices. We present a general framework for heterogeneous DNN compilation, offering automatic partitioning and device mapping. Our scheduler integrates both an exact solver, through a mixed integer linear programming (MILP) formulation, and a modularity-based heuristic for scalability. Furthermore, we propose a theoretical lower bound formula for the optimal solution, which enables the assessment of the heuristic solutions' quality. We evaluate our scheduler in optimizing both conventional DNNs and randomly-wired neural networks, subject to latency and throughput constraints, on a heterogeneous system comprised of a CPU and two distinct GPUs. Compared to na"ively running DNNs on the fastest GPU, he proposed framework can achieve more than 3× times lower latency and up to 2.9× higher throughput by automatically leveraging both data and model parallelism to deploy DNNs on our sample heterogeneous server node. Moreover, our modularity-based "splitting" heuristic improves the solution runtime up to 395× without noticeably sacrificing solution quality compared to an exact MILP solution, and outperforms all other heuristics by 30-60% solution quality. Finally, our case study shows how we can extend our framework to schedule large language models across multiple heterogeneous servers by exploiting symmetry in the hardware setup. Our code can be easily plugged in to existing frameworks, and is available at https://github.com/abdelfattah-lab/diviml.