A New Device Placement Approach with Dual Graph Mamba Networks and Proximal Policy Optimization
摘要
In recent years, foundation models have achieved state-of-the-art performance, but the computational capabilities of GPUs have not kept up, resulting in bottlenecks. While data parallelism, model parallelism, and pipeline parallelism are common in large-scale clusters, automatic parallelism within a single node remains NP-hard. Reinforcement learning methods show promise but are limited by the capabilities of their modules. To tackle these challenges, we propose a novel device placement method (PPO-DGMA) that combines Proximal Policy Optimization (PPO) with a new dual-branch graph embedding approach(DGMA), featuring a decoupled graph neural network (DGM) and a Graph MAMBA (GMA) branch. The PPO algorithm captures subtle environmental changes for enhanced learning and adaptability. Howerver, our dual-branch structure have the capability to improve the extraction of computation graph’s features, and the GMA branch enhances long-distance dependencies, while the DGM branch generates robust node representations by integrating node features, topology graphs, and a semantic graph constructed based on operator attributes. Experimental results show that our method has enhanced execution time by an average of 16.58% over Placeto, 9.65% over GraphSAGE, 3.59% over P-GNN, and 2.81% over CP-GNNAK, while achieving significant reductions in computational time-32.53 \(\times \) faster than Placeto, 37.18 \(\times \) faster than GraphSAGE, 31.34 \(\times \) faster than P-GNN, and 8.68 \(\times \) faster than CP-GNNAK.