Auto-MAP: A DQN Framework for Exploring Distributed Execution Plans for DNN Workloads

2020/07/08 by Siyu Wang, Wang, Siyu, Yi Rong +15
Computer Science · Engineering · #Advanced Neural Network Applications #Artificial Intelligence (cs.AI) #Distributed #FOS: Computer and information sciences #Ferroelectric and Negative Capacitance Devices #Parallel #Parallel Computing and Optimization Techniques #and Cluster Computing (cs.DC)

paper · pdf · doi:10.48550/arxiv.2007.04069

openalex publication_date 2020/07/08 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

The last decade has witnessed growth in the computational requirements for training deep neural networks. Current approaches (e.g., data/model parallelism, pipeline parallelism) parallelize training tasks onto multiple devices. However, these approaches always rely on specific deep learning frameworks and requires elaborate manual design, which make it difficult to maintain and share between different type of models. In this paper, we propose Auto-MAP, a framework for exploring distributed execution plans for DNN workloads, which can automatically discovering fast parallelization strategies through reinforcement learning on IR level of deep learning models. Efficient exploration remains a major challenge for reinforcement learning. We leverage DQN with task-specific pruning strategies to help efficiently explore the search space including optimized strategies. Our evaluation shows that Auto-MAP can find the optimal solution in two hours, while achieving better throughput on several NLP and convolution models.

Citations

• Adam: A Method for Stochastic Optimization
• Very Deep Convolutional Networks for Large-Scale Image Recognition
• Playing Atari with Deep Reinforcement Learning
• Dueling Network Architectures for Deep Reinforcement Learning
• MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems
• Prioritized Experience Replay
• Noisy Networks for Exploration
• PipeDream: Fast and Efficient Pipeline Parallel DNN Training
• Rainbow: Combining Improvements in Deep Reinforcement Learning
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
• GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding

Related

• Beyond Data and Model Parallelism for Deep Neural Networks
• Tessel: Boosting Distributed Execution of Large DNN Models via Flexible Schedule Search
• Demystifying Parallel and Distributed Deep Learning
• ROAM: memory-efficient large DNN training via optimized operator ordering and memory layout
• Demystifying Parallel and Distributed Deep Learning: An In-Depth Concurrency Analysis