Device placement using Laplacian PCA and graph attention networks

Abstract The exponential growth in data and parameters in modern neural networks has created the need to distribute these models across multiple devices for efficient training, resulting in the device placement problem. Existing graph encoding approaches for device placement suffer from low efficiency when searching for optimal parallel strategies, primarily due to suboptimal positional information retrieval. To address these challenges, we propose the Laplacian Principal Component Analysis-graph attention networks (LPCA-GAT) model. Firstly, we employ GAT to capture complex relationships between nodes and generate node encodings. Secondly, we leverage LPCA on the graph Laplacian matrix to extract crucial low-dimensional positional information. Finally, by integrating these two components, we obtain the final node encodings. This enhances the representation capability of nodes within the graph, enabling efficient device placement. The experimental results demonstrate that LPCA-GAT achieves superior device placement results, specifically accelerating execution and computation time by 13.24% and 96.38%, respectively, leading to significant improvements in both operational efficiency and performance.