Optimalisatie van Geheugentoewijzing met Diep Reinforcement Learning

Optimizing Memory Mapping Using Deep Reinforcement Learning

May 11, 2023

Auteurs: 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

Samenvatting

Resource scheduling en allocatie zijn cruciale componenten van veel systemen met grote impact, variërend van congestiebeheer tot cloudcomputing. Het vinden van meer optimale oplossingen voor deze problemen heeft vaak een aanzienlijke invloed op het besparen van middelen en tijd, het verminderen van slijtage aan apparatuur, en zelfs het potentieel verbeteren van de CO2-uitstoot. In dit artikel richten we ons op een specifiek geval van een planningsprobleem, namelijk het geheugenmappingsprobleem dat optreedt tijdens de compilatie van machine learning-programma's: het toewijzen van tensors aan verschillende geheugenlagen om de uitvoeringstijd te optimaliseren. We introduceren een aanpak voor het oplossen van het geheugenmappingsprobleem met behulp van Reinforcement Learning (RL). RL is een oplossingsparadigma dat goed geschikt is voor sequentiële beslissingsproblemen die zich lenen voor planning, en combinatorische zoekruimtes met hoogdimensionale data-inputs. We formuleren het probleem als een single-player game, die we de mallocGame noemen, zodanig dat trajecten met een hoge beloning in het spel overeenkomen met efficiënte geheugenmappingen op de doelhardware. We introduceren ook een Reinforcement Learning-agent, mallocMuZero, en laten zien dat deze in staat is om dit spel te spelen om nieuwe en verbeterde geheugenmappingoplossingen te ontdekken die leiden tot snellere uitvoeringstijden bij echte ML-workloads op ML-accelerators. We vergelijken de prestaties van mallocMuZero met de standaard solver die wordt gebruikt door de Accelerated Linear Algebra (XLA)-compiler op een benchmark van realistische ML-workloads. Daarnaast tonen we aan dat mallocMuZero in staat is om de uitvoeringstijd van het recent gepubliceerde AlphaTensor matrixvermenigvuldigingsmodel te verbeteren.

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.