심층 강화 학습을 활용한 메모리 매핑 최적화

초록

자원 스케줄링 및 할당은 혼잡 제어부터 클라우드 컴퓨팅에 이르기까지 다양한 고성능 시스템의 핵심 구성 요소입니다. 이러한 문제에 대해 더 최적의 해결책을 찾는 것은 자원과 시간 절약, 장치 마모 감소, 심지어 탄소 배출 개선에도 상당한 영향을 미칠 수 있습니다. 본 논문에서는 머신러닝 프로그램 컴파일 중 발생하는 메모리 매핑 문제, 즉 텐서를 서로 다른 메모리 계층에 매핑하여 실행 시간을 최적화하는 특정 스케줄링 문제에 초점을 맞춥니다. 우리는 강화 학습(Reinforcement Learning, RL)을 사용하여 메모리 매핑 문제를 해결하는 접근 방식을 소개합니다. RL은 계획 가능한 순차적 의사결정 문제와 고차원 데이터 입력을 가진 조합적 탐색 공간에 적합한 해결 패러다임입니다. 우리는 이 문제를 단일 플레이어 게임으로 공식화하고, 이를 mallocGame이라고 명명했습니다. 이 게임에서 높은 보상을 얻는 궤적은 대상 하드웨어에서 효율적인 메모리 매핑에 해당합니다. 또한, 우리는 강화 학습 에이전트인 mallocMuZero를 소개하고, 이 에이전트가 이 게임을 통해 새로운 개선된 메모리 매핑 솔루션을 발견하여 실제 ML 가속기에서의 ML 워크로드 실행 시간을 단축할 수 있음을 보여줍니다. 우리는 mallocMuZero의 성능을 Accelerated Linear Algebra (XLA) 컴파일러에서 사용되는 기본 솔버와 현실적인 ML 워크로드 벤치마크에서 비교합니다. 추가적으로, 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