A Structure-Aware Framework for Learning Device Placements on Computation Graphs
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.
Code (1)
Tasks
CPUgraph partitioningGraph Representation Learningreinforcement-learningRepresentation LearningMethods 이 논문이 사용한 방법론
Similar Papers 제목 키워드 기반
Learning Generalizable Device Placement Algorithms for Distributed Machine Learning
We present Placeto, a reinforcement learning (RL) approach to efficiently find device placements for distributed neural network training. Unlike prior approaches that only find a device placement for a specific computat…
BIG-bench Machine LearningReinforcement LearningReinforcement Learning (RL)Placeto: Learning Generalizable Device Placement Algorithms for Distributed Machine Learning
We present Placeto, a reinforcement learning (RL) approach to efficiently find device placements for distributed neural network training. Unlike prior approaches that only find a device placement for a specific computati…
BIG-bench Machine LearningReinforcement LearningReinforcement Learning (RL)A Hierarchical Model for Device Placement
We introduce a hierarchical model for efficient placement of computational graphs onto hardware devices, especially in heterogeneous environments with a mixture of CPUs, GPUs, and other computational devices. Our method …
Deep Reinforcement LearningMachine TranslationmodelReinforcement Learning+2Device Placement Optimization with Reinforcement Learning
The past few years have witnessed a growth in size and computational requirements for training and inference with neural networks. Currently, a common approach to address these requirements is to use a heterogeneous dist…
Language ModelingLanguage ModellingMachine Translationreinforcement-learning+3Baechi: Fast Device Placement of Machine Learning Graphs
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. …