# A Transferable Approach for Partitioning Machine Learning Models on Multi-Chip-Modules

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 22 of its 41 references that are in the map.
- First posted on arXiv 7 Dec 2021.

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

- [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)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Conformer: Convolution-augmented Transformer for Speech Recognition](https://mlsanity.com/paper/2005.08100/) (2020)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Exploring the Limits of Weakly Supervised Pretraining](https://mlsanity.com/paper/1805.00932/) (2018)
- [OptNet: Differentiable Optimization as a Layer in Neural Networks](https://mlsanity.com/paper/1703.00443/) (2017)
- [Deep Learning Scaling is Predictable, Empirically](https://mlsanity.com/paper/1712.00409/) (2017)
- [Deep Speech 2: End-to-End Speech Recognition in English and Mandarin](https://mlsanity.com/paper/1512.02595/) (2015)
- [NAS-FPN: Learning Scalable Feature Pyramid Architecture for Object Detection](https://mlsanity.com/paper/1904.07392/) (2019)
- [Pushing the Limits of Semi-Supervised Learning for Automatic Speech Recognition](https://mlsanity.com/paper/2010.10504/) (2020)
- [Exploring the Limits of Language Modeling](https://mlsanity.com/paper/1602.02410/) (2016)
- [SATNet: Bridging deep learning and logical reasoning using a differentiable satisfiability solver](https://mlsanity.com/paper/1905.12149/) (2019)
- [Mesh-TensorFlow: Deep Learning for Supercomputers](https://mlsanity.com/paper/1811.02084/) (2018)
- [Device Placement Optimization with Reinforcement Learning](https://mlsanity.com/paper/1706.04972/) (2017)
- [Transferable Graph Optimizers for ML Compilers](https://mlsanity.com/paper/2010.12438/) (2020)
- [Placeto: Learning Generalizable Device Placement Algorithms for Distributed Machine Learning](https://mlsanity.com/paper/1906.08879/) (2019)
- [IOS: Inter-Operator Scheduler for CNN Acceleration](https://mlsanity.com/paper/2011.01302/) (2020)

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

- [Automatic Discovery of Composite SPMD Partitioning Strategies in PartIR](https://mlsanity.com/paper/2210.06352/) (2022)
- [Practical Performance Guarantees for Pipelined DNN Inference](https://mlsanity.com/paper/2311.03703/) (2023)

## Abstract (arXiv)

Multi-Chip-Modules (MCMs) reduce the design and fabrication cost of machine learning (ML) accelerators while delivering performance and energy efficiency on par with a monolithic large chip. However, ML compilers targeting MCMs need to solve complex optimization problems optimally and efficiently to achieve this high performance. One such problem is the multi-chip partitioning problem where compilers determine the optimal partitioning and placement of operations in tensor computation graphs on chiplets in MCMs. Partitioning ML graphs for MCMs is particularly hard as the search space grows exponentially with the number of chiplets available and the number of nodes in the neural network. Furthermore, the constraints imposed by the underlying hardware produce a search space where valid solutions are extremely sparse. In this paper, we present a strategy using a deep reinforcement learning (RL) framework to emit a possibly invalid candidate partition that is then corrected by a constraint solver. Using the constraint solver ensures that RL encounters valid solutions in the sparse space frequently enough to converge with fewer samples as compared to non-learned strategies. The architectural choices we make for the policy network allow us to generalize across different ML graphs. Our evaluation of a production-scale model, BERT, on real hardware reveals that the partitioning generated using RL policy achieves 6.11% and 5.85% higher throughput than random search and simulated annealing. In addition, fine-tuning the pre-trained RL policy reduces the search time from 3 hours to only 9 minutes, while achieving the same throughput as training RL policy from scratch.