Optimización del Mapeo de Memoria mediante Aprendizaje por Refuerzo Profundo

Resumen

La programación y asignación de recursos es un componente crítico en muchos sistemas de alto impacto, desde el control de congestión hasta la computación en la nube. Encontrar soluciones más óptimas para estos problemas suele tener un impacto significativo en el ahorro de recursos y tiempo, reducir el desgaste de los dispositivos e incluso mejorar potencialmente las emisiones de carbono. En este artículo, nos centramos en un caso específico de un problema de programación: el problema de mapeo de memoria que ocurre durante la compilación de programas de aprendizaje automático. Específicamente, nos referimos al mapeo de tensores en diferentes capas de memoria para optimizar el tiempo de ejecución. Presentamos un enfoque para resolver el problema de mapeo de memoria utilizando Aprendizaje por Refuerzo (RL, por sus siglas en inglés). RL es un paradigma de solución adecuado para problemas de toma de decisiones secuenciales que se prestan a la planificación y para espacios de búsqueda combinatorios con entradas de datos de alta dimensionalidad. Formulamos el problema como un juego de un solo jugador, al que llamamos mallocGame, de modo que las trayectorias de alta recompensa del juego correspondan a mapeos de memoria eficientes en el hardware objetivo. También presentamos un agente de Aprendizaje por Refuerzo, mallocMuZero, y demostramos que es capaz de jugar este juego para descubrir nuevas y mejoradas soluciones de mapeo de memoria que conducen a tiempos de ejecución más rápidos en cargas de trabajo reales de aprendizaje automático en aceleradores de ML. Comparamos el rendimiento de mallocMuZero con el solucionador predeterminado utilizado por el compilador de Álgebra Lineal Acelerada (XLA) en un conjunto de pruebas de cargas de trabajo realistas de ML. Además, demostramos que mallocMuZero es capaz de mejorar el tiempo de ejecución del modelo de multiplicación de matrices AlphaTensor, publicado recientemente.

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.

Optimización del Mapeo de Memoria mediante Aprendizaje por Refuerzo Profundo

Optimizing Memory Mapping Using Deep Reinforcement Learning

Resumen

Support