Optimizing Memory Mapping Using Deep Reinforcement Learning

by   Pengming Wang, et al.

Resource scheduling and allocation is a critical component of many high impact systems ranging from congestion control to cloud computing. Finding more optimal solutions to these problems often has significant impact on resource and time savings, reducing device wear-and-tear, and even potentially improving carbon emissions. In this paper, we focus on a specific instance of a scheduling problem, namely the memory mapping problem that occurs during compilation of machine learning programs: That is, mapping tensors to different memory layers to optimize execution time. We introduce an approach for solving the memory mapping problem using Reinforcement Learning. RL is a solution paradigm well-suited for sequential decision making problems that are amenable to planning, and combinatorial search spaces with high-dimensional data inputs. We formulate the problem as a single-player game, which we call the mallocGame, such that high-reward trajectories of the game correspond to efficient memory mappings on the target hardware. We also introduce a Reinforcement Learning agent, mallocMuZero, and show that it is capable of playing this game to discover new and improved memory mapping solutions that lead to faster execution times on real ML workloads on ML accelerators. We compare the performance of mallocMuZero to the default solver used by the Accelerated Linear Algebra (XLA) compiler on a benchmark of realistic ML workloads. In addition, we show that mallocMuZero is capable of improving the execution time of the recently published AlphaTensor matrix multiplication model.


page 4

page 18

page 24


Deep Reinforcement Learning-based Methods for Resource Scheduling in Cloud Computing: A Review and Future Directions

As the quantity and complexity of information processed by software syst...

Dyna-H: a heuristic planning reinforcement learning algorithm applied to role-playing-game strategy decision systems

In a Role-Playing Game, finding optimal trajectories is one of the most ...

Task Placement and Resource Allocation for Edge Machine Learning: A GNN-based Multi-Agent Reinforcement Learning Paradigm

Machine learning (ML) tasks are one of the major workloads in today's ed...

Quality of service based radar resource management using deep reinforcement learning

An intelligent radar resource management is an essential milestone in th...

VMAgent: Scheduling Simulator for Reinforcement Learning

A novel simulator called VMAgent is introduced to help RL researchers be...

Optimizing Memory Placement using Evolutionary Graph Reinforcement Learning

As modern neural networks have grown to billions of parameters, meeting ...

Reinforcement Learning (RL) Augmented Cold Start Frequency Reduction in Serverless Computing

Function-as-a-Service is a cloud computing paradigm offering an event-dr...

Please sign up or login with your details

Forgot password? Click here to reset