A novel device placement approach based on position-aware subgraph neural networks

Coping with the growing demand for data and parameters in complex neural network (NN) models of contemporary times typically involves distributing them across multiple devices, giving rise to the device placement problem. Several widely adopted solutions to this challenge leverage graph embedding schemes. However, the current graph embedding models employed for device placement suffer from several limitations, such as being highly time-consuming in the search for optimal parallelization strategies and yielding suboptimal parallel results. In this paper, we develop a novel graph embedding approach CP-GNNAK for efficient device placement. CP-GNNAK mainly depends on an efficient subgraph-based node embedding strategy, namely GNN as kernel, and a novel cosine phase position embedding scheme. CP-GNNAK goes beyond encoding only the immediate neighbors by deriving the representation of each node from the encoding of a surrounding subgraph. This expanded approach significantly enhances the efficiency of device placement through the integration of subgraph sampling strategies. Additionally, the cosine phase position embedding not only considers the positional information between nodes but also captures the directional information underlying the nodes. This adds richness to the node embeddings and improves the performance of device placement. The experimental results clearly indicate that the CP-GNNAK model achieves the state-of-the-art device placement results. Specifically, the experimental results reveal a remarkable 13.77% improvement in runtime when employing CP-GNNAK compared to Placeto. CP-GNNAK also outperforms GraphSAGE by 7.36% and P-GNN by 0.94%. Turning to computation time, CP-GNNAK exhibits a substantial speedup of 3.40× over Placeto, 3.87× over GraphSAGE, and 3.20× over P-GNN.