Optimizing Memory Mapping Using Deep Reinforcement Learning
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.
6461cce85ea1bc9a9da954a0
Pengming Wang
6461cce85ea1bc9a9da954a1
Mikita Sazanovich
6461cce85ea1bc9a9da954a2
Berkin Ilbeyi
6461cce85ea1bc9a9da954a3
Phitchaya Mangpo Phothilimthana
6461cce85ea1bc9a9da954a4
Manish Purohit
admin_assigned
2023-05-26T09:41:42.213Z
6122fde3d09e016ab5187a5f
/avatars/366e14ff8768e8deaee0ff7b21b9f76a.svg
manish purohit
user
manishpurohit
6461cce85ea1bc9a9da954a5
Han Yang Tay
6461cce85ea1bc9a9da954a6
Ngân Vũ
6461cce85ea1bc9a9da954a7
Miaosen Wang
6461cce85ea1bc9a9da954a8
Cosmin Paduraru
6461cce85ea1bc9a9da954a9
Edouard Leurent
admin_assigned
2023-05-26T09:42:32.843Z
633e093315033639f991760f
/avatars/0339ec428a8d1fb961cb0b095e0839bd.svg
Edouard Leurent
user
eleurent
6461cce85ea1bc9a9da954aa
Anton Zhernov
6461cce85ea1bc9a9da954ab
Julian Schrittwieser
6461cce85ea1bc9a9da954ac
Thomas Hubert
6461cce85ea1bc9a9da954ad
Robert Tung
6461cce85ea1bc9a9da954ae
Paula Kurylowicz
6461cce85ea1bc9a9da954af
Kieran Milan
6461cce85ea1bc9a9da954b0
Oriol Vinyals
6461cce85ea1bc9a9da954b1
Daniel J. Mankowitz
Optimizing Memory Mapping Using Deep Reinforcement Learning
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.
A reinforcement learning approach, mallocMuZero, is introduced to optimize memory mapping during ML program compilation, achieving faster execution times on ML accelerators compared to default solvers.
Reinforcement Learning RL memory mapping problem mallocGame mallocMuZero Accelerated Linear Algebra (XLA) AlphaTensor matrix multiplication ML workloads ML accelerators execution time