Оптимизация отображения памяти с использованием глубокого обучения с подкреплением

Аннотация

Планирование и распределение ресурсов являются критически важными компонентами многих высокоэффективных систем, начиная от управления перегрузками и заканчивая облачными вычислениями. Поиск более оптимальных решений для этих задач часто приводит к значительной экономии ресурсов и времени, снижению износа оборудования и даже потенциальному улучшению показателей выбросов углерода. В данной статье мы сосредоточимся на конкретном примере задачи планирования, а именно на проблеме распределения памяти, возникающей при компиляции программ машинного обучения: то есть на задаче отображения тензоров на различные уровни памяти для оптимизации времени выполнения. Мы представляем подход к решению проблемы распределения памяти с использованием обучения с подкреплением (Reinforcement Learning, RL). RL — это парадигма решения, хорошо подходящая для задач последовательного принятия решений, которые допускают планирование, а также для комбинаторных пространств поиска с высокоразмерными входными данными. Мы формулируем задачу как однопользовательскую игру, которую называем mallocGame, таким образом, что траектории с высоким вознаграждением в игре соответствуют эффективным отображениям памяти на целевом оборудовании. Мы также представляем агента обучения с подкреплением, mallocMuZero, и показываем, что он способен играть в эту игру для обнаружения новых и улучшенных решений по распределению памяти, которые приводят к более быстрому выполнению реальных задач машинного обучения на ML-акселераторах. Мы сравниваем производительность mallocMuZero с решателем по умолчанию, используемым компилятором Accelerated Linear Algebra (XLA), на наборе реалистичных задач машинного обучения. Кроме того, мы демонстрируем, что mallocMuZero способен улучшить время выполнения недавно опубликованной модели умножения матриц AlphaTensor.

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.

Оптимизация отображения памяти с использованием глубокого обучения с подкреплением

Optimizing Memory Mapping Using Deep Reinforcement Learning

Аннотация

Support