Optimierung der Speicherzuordnung durch Deep Reinforcement Learning

Zusammenfassung

Ressourcenplanung und -zuweisung ist ein entscheidender Bestandteil vieler hochwirksamer Systeme, die von der Überlastungskontrolle bis hin zum Cloud Computing reichen. Die Suche nach optimaleren Lösungen für diese Probleme hat oft erhebliche Auswirkungen auf die Einsparung von Ressourcen und Zeit, die Reduzierung von Geräteverschleiß und sogar die potenzielle Verbesserung der Kohlenstoffemissionen. In diesem Artikel konzentrieren wir uns auf eine spezifische Instanz eines Planungsproblems, nämlich das Speicherzuordnungsproblem, das während der Kompilierung von maschinellen Lernprogrammen auftritt: das heißt, die Zuordnung von Tensoren zu verschiedenen Speicherebenen, um die Ausführungszeit zu optimieren. Wir stellen einen Ansatz zur Lösung des Speicherzuordnungsproblems mittels Reinforcement Learning (RL) vor. RL ist ein Lösungsparadigma, das sich gut für sequenzielle Entscheidungsprobleme eignet, die sich für die Planung eignen, sowie für kombinatorische Suchräume mit hochdimensionalen Dateneingaben. Wir formulieren das Problem als Einzelspieler-Spiel, das wir mallocGame nennen, sodass hochbewertete Spielverläufe effiziente Speicherzuordnungen auf der Zielhardware entsprechen. Wir führen auch einen Reinforcement-Learning-Agenten, mallocMuZero, ein und zeigen, dass dieser in der Lage ist, dieses Spiel zu spielen, um neue und verbesserte Speicherzuordnungslösungen zu entdecken, die zu schnelleren Ausführungszeiten bei realen ML-Workloads auf ML-Beschleunigern führen. Wir vergleichen die Leistung von mallocMuZero mit dem Standardlöser, der vom Accelerated Linear Algebra (XLA)-Compiler verwendet wird, anhand eines Benchmarks realistischer ML-Workloads. Darüber hinaus zeigen wir, dass mallocMuZero in der Lage ist, die Ausführungszeit des kürzlich veröffentlichten AlphaTensor-Matrixmultiplikationsmodells zu verbessern.

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.

Optimierung der Speicherzuordnung durch Deep Reinforcement Learning

Optimizing Memory Mapping Using Deep Reinforcement Learning

Zusammenfassung

Support