Optimizing Memory Map** Using Deep Reinforcement Learning
Authors:
Pengming Wang,
Mikita Sazanovich,
Berkin Ilbeyi,
Phitchaya Mangpo Phothilimthana,
Manish Purohit,
Han Yang Tay,
Ngân Vũ,
Miaosen Wang,
Cosmin Paduraru,
Edouard Leurent,
Anton Zhernov,
Po-Sen Huang,
Julian Schrittwieser,
Thomas Hubert,
Robert Tung,
Paula Kurylowicz,
Kieran Milan,
Oriol Vinyals,
Daniel J. Mankowitz
Abstract:
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, n…
▽ More
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 map** problem that occurs during compilation of machine learning programs: That is, map** tensors to different memory layers to optimize execution time.
We introduce an approach for solving the memory map** 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 map**s 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 map** 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.
△ Less
Submitted 17 October, 2023; v1 submitted 11 May, 2023;
originally announced May 2023.
MuZero with Self-competition for Rate Control in VP9 Video Compression
Authors:
Amol Mandhane,
Anton Zhernov,
Maribeth Rauh,
Chenjie Gu,
Miaosen Wang,
Flora Xue,
Wendy Shang,
Derek Pang,
Rene Claus,
Ching-Han Chiang,
Cheng Chen,
**gning Han,
Angie Chen,
Daniel J. Mankowitz,
Jackson Broshear,
Julian Schrittwieser,
Thomas Hubert,
Oriol Vinyals,
Timothy Mann
Abstract:
Video streaming usage has seen a significant rise as entertainment, education, and business increasingly rely on online video. Optimizing video compression has the potential to increase access and quality of content to users, and reduce energy use and costs overall. In this paper, we present an application of the MuZero algorithm to the challenge of video compression. Specifically, we target the p…
▽ More
Video streaming usage has seen a significant rise as entertainment, education, and business increasingly rely on online video. Optimizing video compression has the potential to increase access and quality of content to users, and reduce energy use and costs overall. In this paper, we present an application of the MuZero algorithm to the challenge of video compression. Specifically, we target the problem of learning a rate control policy to select the quantization parameters (QP) in the encoding process of libvpx, an open source VP9 video compression library widely used by popular video-on-demand (VOD) services. We treat this as a sequential decision making problem to maximize the video quality with an episodic constraint imposed by the target bitrate. Notably, we introduce a novel self-competition based reward mechanism to solve constrained RL with variable constraint satisfaction difficulty, which is challenging for existing constrained RL methods. We demonstrate that the MuZero-based rate control achieves an average 6.28% reduction in size of the compressed videos for the same delivered video quality level (measured as PSNR BD-rate) compared to libvpx's two-pass VBR rate control policy, while having better constraint satisfaction behavior.
△ Less
Submitted 14 February, 2022;
originally announced February 2022.
