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.

错误:搜索内容不能为空,请输入英文关键词
错误:关键词超出字数限制,请精简
高级检索

A New Device Placement Approach with Dual Graph Mamba Networks and Proximal Policy Optimization

  • Meng Han,
  • Yan Zeng,
  • Hao Shu,
  • Xiaofei Lu,
  • Jilin Zhang,
  • Yongjian Ren,
  • Wangli Hao

摘要

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.