OpenGraphGym-MG: Using Reinforcement Learning to Solve Large Graph Optimization Problems on MultiGPU Systems

2021/05/18 by Weijian Zheng, Zheng, Weijian, Dali Wang +3
Computer Science · Engineering · #Advanced Graph Neural Networks #Artificial intelligence #Caching and Content Delivery #Computer science #Distributed #Embedding #FOS: Computer and information sciences #Ferroelectric and Negative Capacitance Devices #Graph #Graph embedding #Inference #Parallel #Parallel computing #Reinforcement learning #Scalability #Supercomputer #Theoretical computer science #and Cluster Computing (cs.DC) #cs.DC

paper · pdf · doi:10.48550/arxiv.2105.08764

published in arXiv (Cornell University) (Cornell University)

openalex publication_date 2021/05/18 · arxiv created 2021/06/24 · arxiv updated 2021/06/25 · openalex created_date 2025/10/10 · openalex updated_date 2026/08/08

Abstract

Large scale graph optimization problems arise in many fields. This paper presents an extensible, high performance framework (named OpenGraphGym-MG) that uses deep reinforcement learning and graph embedding to solve large graph optimization problems with multiple GPUs. The paper uses a common RL algorithm (deep Q-learning) and a representative graph embedding (structure2vec) to demonstrate the extensibility of the framework and, most importantly, to illustrate the novel optimization techniques, such as spatial parallelism, graph-level and node-level batched processing, distributed sparse graph storage, efficient parallel RL training and inference algorithms, repeated gradient descent iterations, and adaptive multiple-node selections. This study performs a comprehensive performance analysis on parallel efficiency and memory cost that proves the parallel RL training and inference algorithms are efficient and highly scalable on a number of GPUs. This study also conducts a range of large graph experiments, with both generated graphs (over 30 million edges) and real-world graphs, using a single compute node (with six GPUs) of the Summit supercomputer. Good scalability in both RL training and inference is achieved: as the number of GPUs increases from one to six, the time of a single step of RL training and a single step of RL inference on large graphs with more than 30 million edges, is reduced from 316.4s to 54.5s, and 23.8s to 3.4s, respectively. The research results on a single node lay out a solid foundation for the future work to address graph optimization problems with a large number of GPUs across multiple nodes in the Summit.

Citations

• Adam: A Method for Stochastic Optimization
• word2vec Explained: deriving Mikolov et al.'s negative-sampling word-embedding method
• Neural Network for Graphs: A Contextual Constructive Approach
• Discriminative Embeddings of Latent Variable Models for Structured Data
• Neural Message Passing for Quantum Chemistry
• Learning Combinatorial Optimization Algorithms over Graphs
• Inductive Representation Learning on Large Graphs
• Fast Graph Representation Learning with PyTorch Geometric
• Learning Heuristics over Large Graphs via Deep Reinforcement Learning
• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• Accelerating Reinforcement Learning through GPU Atari Emulation
• Acme: A Research Framework for Distributed Reinforcement Learning
• DistDGL: Distributed Graph Neural Network Training for Billion-Scale Graphs
• Ecole: A Gym-like Library for Machine Learning in Combinatorial Optimization Solvers
• A Comprehensive Survey on Graph Neural Networks

Related

• Scalable Graph Embedding LearningOn A Single GPU
• MG-GCN: Scalable Multi-GPU GCN Training Framework
• Multi-GPU Graph Analytics
• BGL: GPU-Efficient GNN Training by Optimizing Graph Data I/O and Preprocessing
• Plexus: Taming Billion-edge Graphs with 3D Parallel Full-graph GNN Training