利用深度強化學習優化記憶映射

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.