# A Structure-Aware Framework for Learning Device Placements on Computation Graphs

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 25 of its 38 references that are in the map.
- First posted on arXiv 23 May 2024.

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

- [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)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [Rethinking the Inception Architecture for Computer Vision](https://mlsanity.com/paper/1512.00567/) (2015)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [Variational Graph Auto-Encoders](https://mlsanity.com/paper/1611.07308/) (2016)
- [Hierarchical Graph Representation Learning with Differentiable Pooling](https://mlsanity.com/paper/1806.08804/) (2018)
- [Device Placement Optimization with Reinforcement Learning](https://mlsanity.com/paper/1706.04972/) (2017)
- [Position: Graph Foundation Models are Already Here](https://mlsanity.com/paper/2402.02216/) (2024)
- [Placeto: Learning Generalizable Device Placement Algorithms for Distributed Machine Learning](https://mlsanity.com/paper/1906.08879/) (2019)
- [Enhancing Computation Efficiency in Large Language Models through Weight and Activation Quantization](https://mlsanity.com/paper/2311.05161/) (2023)
- [Neuro-Inspired Information-Theoretic Hierarchical Perception for Multimodal Learning](https://mlsanity.com/paper/2404.09403/) (2024)
- [OMPGPT: A Generative Pre-trained Transformer Model for OpenMP](https://mlsanity.com/paper/2401.16445/) (2024)
- [A Comprehensive Survey on Distributed Training of Graph Neural Networks](https://mlsanity.com/paper/2211.05368/) (2022)
- [Structure Guided Prompt: Instructing Large Language Model in Multi-Step Reasoning by Exploring Graph Structure of the Text](https://mlsanity.com/paper/2402.13415/) (2024)
- [Leveraging Large Language Models to Detect Influence Campaigns in Social Media](https://mlsanity.com/paper/2311.07816/) (2023)
- [GDP: Generalized Device Placement for Dataflow Graphs](https://mlsanity.com/paper/1910.01578/) (2019)
- [Neuron-based Multifractal Analysis of Neuron Interaction Dynamics in Large Models](https://mlsanity.com/paper/2402.09099/) (2024)
- [MPIrigen: MPI Code Generation through Domain-Specific Language Models](https://mlsanity.com/paper/2402.09126/) (2024)
- [Graph Parsing Networks](https://mlsanity.com/paper/2402.14393/) (2024)
- [End-to-end Mapping in Heterogeneous Systems Using Graph Representation Learning](https://mlsanity.com/paper/2204.11981/) (2022)

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

- [POET: Power-Oriented Evolutionary Tuning for LLM-Based RTL PPA Optimization](https://mlsanity.com/paper/2603.19333/) (2026)
- [Resource Heterogeneity-Aware and Utilization-Enhanced Scheduling for Deep Learning Clusters](https://mlsanity.com/paper/2503.10918/) (2025)
- [DOPPLER: Dual-Policy Learning for Device Assignment in Asynchronous Dataflow Graphs](https://mlsanity.com/paper/2505.23131/) (2025)

## Abstract (arXiv)

Computation graphs are Directed Acyclic Graphs (DAGs) where the nodes correspond to mathematical operations and are used widely as abstractions in optimizations of neural networks. The device placement problem aims to identify optimal allocations of those nodes to a set of (potentially heterogeneous) devices. Existing approaches rely on two types of architectures known as grouper-placer and encoder-placer, respectively. In this work, we bridge the gap between encoder-placer and grouper-placer techniques and propose a novel framework for the task of device placement, relying on smaller computation graphs extracted from the OpenVINO toolkit. The framework consists of five steps, including graph coarsening, node representation learning and policy optimization. It facilitates end-to-end training and takes into account the DAG nature of the computation graphs. We also propose a model variant, inspired by graph parsing networks and complex network analysis, enabling graph representation learning and jointed, personalized graph partitioning, using an unspecified number of groups. To train the entire framework, we use reinforcement learning using the execution time of the placement as a reward. We demonstrate the flexibility and effectiveness of our approach through multiple experiments with three benchmark models, namely Inception-V3, ResNet, and BERT. The robustness of the proposed framework is also highlighted through an ablation study. The suggested placements improve the inference speed for the benchmark models by up to 58.2% over CPU execution and by up to 60.24% compared to other commonly used baselines.