Ottimizzazione del Mapping della Memoria Utilizzando l'Apprendimento per Rinforzo Profondo

Optimizing Memory Mapping Using Deep Reinforcement Learning

May 11, 2023

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

Abstract

La pianificazione e l'allocazione delle risorse sono componenti critici di molti sistemi ad alto impatto, che vanno dal controllo della congestione al cloud computing. Trovare soluzioni più ottimali per questi problemi spesso ha un impatto significativo sul risparmio di risorse e tempo, riducendo l'usura dei dispositivi e potenzialmente migliorando anche le emissioni di carbonio. In questo articolo, ci concentriamo su un'istanza specifica di un problema di pianificazione, ovvero il problema di mappatura della memoria che si verifica durante la compilazione di programmi di machine learning: cioè, la mappatura dei tensori su diversi livelli di memoria per ottimizzare il tempo di esecuzione. Introduciamo un approccio per risolvere il problema di mappatura della memoria utilizzando il Reinforcement Learning (RL). L'RL è un paradigma di soluzione ben adatto per problemi di decisione sequenziale che si prestano alla pianificazione e per spazi di ricerca combinatoria con input di dati ad alta dimensionalità. Formuliamo il problema come un gioco a giocatore singolo, che chiamiamo mallocGame, in modo che le traiettorie ad alta ricompensa del gioco corrispondano a mappature di memoria efficienti sull'hardware target. Introduciamo anche un agente di Reinforcement Learning, mallocMuZero, e dimostriamo che è in grado di giocare a questo gioco per scoprire nuove e migliori soluzioni di mappatura della memoria che portano a tempi di esecuzione più rapidi su carichi di lavoro ML reali su acceleratori ML. Confrontiamo le prestazioni di mallocMuZero con il risolutore predefinito utilizzato dal compilatore Accelerated Linear Algebra (XLA) su un benchmark di carichi di lavoro ML realistici. Inoltre, dimostriamo che mallocMuZero è in grado di migliorare il tempo di esecuzione del modello di moltiplicazione di matrici AlphaTensor, pubblicato di recente.

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.