利用深度强化学习优化内存映射

Optimizing Memory Mapping Using Deep Reinforcement Learning

May 11, 2023

作者: 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, Julian Schrittwieser, Thomas Hubert, Robert Tung, Paula Kurylowicz, Kieran Milan, Oriol Vinyals, Daniel J. Mankowitz

cs.AI

摘要

资源调度和分配是许多高影响系统的关键组成部分，范围从拥塞控制到云计算。找到这些问题的更优解通常对资源和时间节省、减少设备磨损，甚至潜在地改善碳排放有重大影响。在本文中，我们专注于调度问题的特定实例，即在机器学习程序编译过程中出现的内存映射问题：即将张量映射到不同的内存层以优化执行时间。 我们介绍了一种使用强化学习解决内存映射问题的方法。强化学习是一种适合于可规划的顺序决策问题和具有高维数据输入的组合搜索空间的解决方案范式。我们将问题制定为单人游戏，我们称之为mallocGame，使得游戏的高奖励轨迹对应于目标硬件上的高效内存映射。我们还介绍了一个强化学习代理，mallocMuZero，并展示它能够玩这个游戏，发现新的和改进的内存映射解决方案，从而在ML加速器上的实际ML工作负载上实现更快的执行时间。我们将mallocMuZero的性能与加速线性代数（XLA）编译器使用的默认求解器在一组真实ML工作负载的基准上进行了比较。此外，我们展示了mallocMuZero能够改善最近发布的AlphaTensor矩阵乘法模型的执行时间。

English

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.