# Partitioning Distributed Compute Jobs with Reinforcement Learning and Graph Neural Networks

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

## Key facts

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

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [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)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [Google's Neural Machine Translation System: Bridging the Gap between Human and Machine Translation](https://mlsanity.com/paper/1609.08144/) (2016)
- [SqueezeNet: AlexNet-level accuracy with 50x fewer parameters and <0.5MB model size](https://mlsanity.com/paper/1602.07360/) (2016)
- [Rainbow: Combining Improvements in Deep Reinforcement Learning](https://mlsanity.com/paper/1710.02298/) (2017)
- [Dueling Network Architectures for Deep Reinforcement Learning](https://mlsanity.com/paper/1511.06581/) (2015)
- [Machine Learning for Combinatorial Optimization: a Methodological Tour d'Horizon](https://mlsanity.com/paper/1811.06128/) (2018)
- [Neural Combinatorial Optimization with Reinforcement Learning](https://mlsanity.com/paper/1611.09940/) (2016)
- [Learning Combinatorial Optimization Algorithms over Graphs](https://mlsanity.com/paper/1704.01665/) (2017)
- [Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model](https://mlsanity.com/paper/2201.11990/) (2022)
- [On the Number of Linear Regions of Deep Neural Networks](https://mlsanity.com/paper/1402.1869/) (2014)
- [Quantifying Generalization in Reinforcement Learning](https://mlsanity.com/paper/1812.02341/) (2018)
- [FeUdal Networks for Hierarchical Reinforcement Learning](https://mlsanity.com/paper/1703.01161/) (2017)
- [Exact Combinatorial Optimization with Graph Convolutional Neural Networks](https://mlsanity.com/paper/1906.01629/) (2019)
- [Distributed Prioritized Experience Replay](https://mlsanity.com/paper/1803.00933/) (2018)
- [Google Research Football: A Novel Reinforcement Learning Environment](https://mlsanity.com/paper/1907.11180/) (2019)
- [A Practical Guide to Multi-Objective Reinforcement Learning and Planning](https://mlsanity.com/paper/2103.09568/) (2021)
- [Generalization and Representational Limits of Graph Neural Networks](https://mlsanity.com/paper/2002.06157/) (2020)
- [Understanding Attention and Generalization in Graph Neural Networks](https://mlsanity.com/paper/1905.02850/) (2019)
- [Demystifying Parallel and Distributed Deep Learning: An In-Depth Concurrency Analysis](https://mlsanity.com/paper/1802.09941/) (2018)
- [Acme: A Research Framework for Distributed Reinforcement Learning](https://mlsanity.com/paper/2006.00979/) (2020)
- [Memory-Efficient Pipeline-Parallel DNN Training](https://mlsanity.com/paper/2006.09503/) (2020)
- [Device Placement Optimization with Reinforcement Learning](https://mlsanity.com/paper/1706.04972/) (2017)
- [Improving Generalization in Reinforcement Learning with Mixture Regularization](https://mlsanity.com/paper/2010.10814/) (2020)

## Abstract (arXiv)

From natural language processing to genome sequencing, large-scale machine learning models are bringing advances to a broad range of fields. Many of these models are too large to be trained on a single machine, and instead must be distributed across multiple devices. This has motivated the research of new compute and network systems capable of handling such tasks. In particular, recent work has focused on developing management schemes which decide how to allocate distributed resources such that some overall objective, such as minimising the job completion time (JCT), is optimised. However, such studies omit explicit consideration of how much a job should be distributed, usually assuming that maximum distribution is desirable. In this work, we show that maximum parallelisation is sub-optimal in relation to user-critical metrics such as throughput and blocking rate. To address this, we propose PAC-ML (partitioning for asynchronous computing with machine learning). PAC-ML leverages a graph neural network and reinforcement learning to learn how much to partition computation graphs such that the number of jobs which meet arbitrary user-defined JCT requirements is maximised. In experiments with five real deep learning computation graphs on a recently proposed optical architecture across four user-defined JCT requirement distributions, we demonstrate PAC-ML achieving up to 56.2% lower blocking rates in dynamic job arrival settings than the canonical maximum parallelisation strategy used by most prior works.