Otimização de Mapeamento de Memória Usando Aprendizado por Reforço Profundo

Optimizing Memory Mapping Using Deep Reinforcement Learning

May 11, 2023

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

Resumo

O agendamento e alocação de recursos é um componente crítico de muitos sistemas de alto impacto, desde o controle de congestionamento até a computação em nuvem. Encontrar soluções mais ótimas para esses problemas frequentemente tem um impacto significativo na economia de recursos e tempo, na redução do desgaste de dispositivos e até mesmo na potencial melhoria das emissões de carbono. Neste artigo, focamos em uma instância específica de um problema de agendamento, especificamente o problema de mapeamento de memória que ocorre durante a compilação de programas de aprendizado de máquina: ou seja, o mapeamento de tensores para diferentes camadas de memória para otimizar o tempo de execução. Introduzimos uma abordagem para resolver o problema de mapeamento de memória usando Aprendizado por Reforço (RL). O RL é um paradigma de solução bem adequado para problemas de tomada de decisão sequencial que são passíveis de planejamento e para espaços de busca combinatórios com entradas de dados de alta dimensionalidade. Formulamos o problema como um jogo de um único jogador, que chamamos de mallocGame, de modo que trajetórias de alta recompensa do jogo correspondam a mapeamentos de memória eficientes no hardware alvo. Também introduzimos um agente de Aprendizado por Reforço, mallocMuZero, e mostramos que ele é capaz de jogar esse jogo para descobrir novas e aprimoradas soluções de mapeamento de memória que levam a tempos de execução mais rápidos em cargas de trabalho reais de ML em aceleradores de ML. Comparamos o desempenho do mallocMuZero com o solucionador padrão usado pelo compilador de Álgebra Linear Acelerada (XLA) em um benchmark de cargas de trabalho realistas de ML. Além disso, mostramos que o mallocMuZero é capaz de melhorar o tempo de execução do modelo de multiplicação de matrizes AlphaTensor, recentemente publicado.

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.