Optimisation du mappage mémoire par apprentissage par renforcement profond

Résumé

La planification et l'allocation des ressources constituent un élément crucial de nombreux systèmes à fort impact, allant du contrôle de la congestion au cloud computing. Trouver des solutions plus optimales à ces problèmes a souvent un impact significatif sur les économies de ressources et de temps, réduit l'usure des appareils, et peut même potentiellement améliorer les émissions de carbone. Dans cet article, nous nous concentrons sur un cas spécifique de problème de planification, à savoir le problème de mappage de mémoire qui survient lors de la compilation de programmes d'apprentissage automatique : c'est-à-dire, le mappage des tenseurs sur différentes couches de mémoire pour optimiser le temps d'exécution. Nous introduisons une approche pour résoudre le problème de mappage de mémoire en utilisant l'apprentissage par renforcement (Reinforcement Learning, RL). Le RL est un paradigme de solution bien adapté aux problèmes de prise de décision séquentielle qui se prêtent à la planification, ainsi qu'aux espaces de recherche combinatoires avec des entrées de données de haute dimension. Nous formulons le problème comme un jeu à un seul joueur, que nous appelons le mallocGame, de telle sorte que les trajectoires à haut rendement du jeu correspondent à des mappages de mémoire efficaces sur le matériel cible. Nous introduisons également un agent d'apprentissage par renforcement, mallocMuZero, et montrons qu'il est capable de jouer à ce jeu pour découvrir de nouvelles solutions de mappage de mémoire améliorées, conduisant à des temps d'exécution plus rapides sur des charges de travail réelles d'apprentissage automatique sur des accélérateurs ML. Nous comparons les performances de mallocMuZero à celles du solveur par défaut utilisé par le compilateur Accelerated Linear Algebra (XLA) sur un benchmark de charges de travail ML réalistes. De plus, nous montrons que mallocMuZero est capable d'améliorer le temps d'exécution du modèle de multiplication matricielle AlphaTensor récemment publié.

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.

Optimisation du mappage mémoire par apprentissage par renforcement profond

Optimizing Memory Mapping Using Deep Reinforcement Learning

Résumé

Support