Optimizing Memory Placement using Evolutionary Graph Reinforcement Learning

1 INTRODUCTION . The proliferation of deep learning ( DL ) has been fueled , in part , by a rapid growth in the size and complexity of deep neural networks ( DNN ) ( Dean et al. , 2012 ; Ying et al. , 2018 ) . This has spurred the rapid development of hardware ( Wang et al. , 2016 ; Jouppi et al. , 2017 ) and software ( Abadi et al. , 2016 ; Paszke et al. , 2018 ; Cyphers et al. , 2018 ) dedicated to deep learning workloads that seek to optimize critical performance metrics like throughput and power efficiency ( Mattson et al. , 2020 ) . Producing compiler optimizations that map the tensors of a neural network ’ s computational graph to the memory units on host hardware is a critical challenge . Since different memory types trade off bandwidth and capacity differently , a sub-optimal mapping could significantly increase latency . For DL inference , the computational graph is static and placement can be pre-planned instead of relying on online cache management ( Zhang et al. , 2020 ; Shi et al. , 2019 ) . However , this is especially challenging with DNNs due to the high dimensional search space . For example , ResNet-50 ( He et al. , 2016 ) has 57 operational layers . Mapping each activation and weight tensor to , for example , three ( DRAM , LLC , and SRAM ) memory caches represents 3 ( 2∗57 ) ≈ 1054 possible decisions . BERT ( Devlin et al. , 2018 ) has 376 operational layers , and a search space of ∼ 10358 . Since optimizing this mapping is intractable with traditional approaches , such as dynamic programming ( Bellman , 1954 ) , current solutions primarily rely on manually-tuned heuristic rules encoded in a compiler . Because of the large search space , prior reinforcement learning ( RL ) algorithms for automating mappings have relied on manually-designed grouping ( Mirhoseini et al. , 2017 ; Addanki et al. , 2018 ) or a learned grouper whose hierarchical structure is domain dependent ( Mirhoseini et al. , 2018 ) . In addition to the extremely large action space , the large number of nodes render the reward sparse and noisy , and thus further unsuitable for gradient-based Deep RL algorithms . This sparsity stems from the fact that an overall performance metric can only be measured after all nodes have been processed . ∗Equal Contribution †Correspondence to : < somdeb.majumdar @ intel.com > In this paper , we present Evolutionary Graph Reinforcement Learning ( EGRL ) , a hybrid approach of evolutionary search with gradient-based learning , that is able to natively search in a high-dimensional space that is orders-of-magnitude larger than previous approaches . EGRL is an extension of CERL ( Khadka et al. , 2019 ) , a population based method for sparse-reward tasks that combines fast policy gradient ( PG ) learning with a stable evolutionary algorithm ( EA ) . Since the action spaces explored in this paper are several orders of magnitude larger than those explored in CERL , we introduce Boltzmann chromosomes - a set of fast , stateless policies that accelerate evolution by providing partially optimized solutions as anchors . This mechanism is necessary to improve the sampleefficiency of the slow EA component for this large action space . Further , we employ a graph neural network ( GNN ) ( Wu et al. , 2020 ; Scarselli et al. , 2008 ) to represent our policy . This allows our agent to natively process computational graphs representing deep learning workloads , enabling generalization over workloads of varying size and connectivity . We demonstrate our solution on the Intel Neural Network Processor for Inference ( NNP-I ) , a deep learning accelerator , to map modern neural networks on one of the three memory hierarchies on the chip . Each memory level in this chip has trade-offs in memory size and bandwidth , as detailed in Wechsler et al . ( 2019 ) . This additionally differentiates our work from prior works such as REGAL ( Paliwal et al. , 2020 ) that assume infinite bandwidth and memory that are not practical on real hardware . Additionally , we consider single-batch inference , an important industry benchmark Mattson et al . ( 2020 ) . While large batch sizes have greater computational efficiency ( e.g. , Boudoukh et al . ( 2020 ) on NNP-I ) , they are sub-optimal for a given inference example due to the latency associated with queuing up a batch . Therefore , single-batch inference is key to many time-critical applications ( Park et al. , 2018 ) where an individual inference query needs to be processed in real-time . Results on ResNet-50 , ResNet-101 ( He et al. , 2016 ) and BERT , show that EGRL significantly outperforms the chipset ’ s native compiler across all workloads , and exceeds the performance of dynamic programming , evolutionary search and policy-gradient approaches . Specifically , the contributions of this work are : 1 . A generalized GNN-based policy that can natively accept a computational graph and produce a corresponding graph representation with the optimal memory maps . This eliminates the need for serialized , layer-dependent representations . 2 . EGRL , a scalable population-based algorithm that can effectively train on sparse and noisy feedback from the host hardware in large search spaces . 3 . An RL agent that trains directly on the hardware , with a feedback mechanism for constraint violation , and thus allowing direct deployment and testing on hardware . 2 RELATED WORK . Optimizing Hardware using Machine Learning : In this work , we study the problem of mapping tensors to memory components on a device . Several recent works have studied the use of machine learning to optimize the execution of computation graphs on hardware -a similar combinatorics problem . Mirhoseini et al . ( 2017 ) designed a policy gradient ( PG ) based Placer policy to map parts of neural models on hardware . However , it relied on a heuristic grouping strategy to significantly reduce the action space . Mirhoseini et al . ( 2018 ) improved this architecture by replacing the heuristic module with a Grouper policy . While this does represent an end-to-end PG approach , it is significantly more complex and hyperparameter-heavy than EGRL ’ s PG network . Specifically , their Placer is an LSTM based Seq-to-Seq model with attention . The hierarchical structure is domain dependent - specific to operation grouping . Mirhoseini et al . ( 2020 ) also applied Deep RL to learn subsystem placement to optimize power , performance and chip area but relied on similar heuristic grouping strategies to significantly reduce the action space seen by the agent . A closely related work is Placeto ( Addanki et al. , 2018 ) where the nodes of a computation graph are sequentially placed on hardware . Similar to EGRL , they also operate on a GNN representation . EGRL primarily differs from Placeto in simultaneously mapping all nodes . We found empirically that a sequential mapping strategy was significantly more sample-inefficient and could not scale to larger workloads . Sequential mapping strategies have the additional disadvantage of not being able to exploit parallelism during policy training . Placeto also adopts the manual grouping strategy from Mirhoseini et al . ( 2018 ) and adds an additional Merge and Collocate heuristic . In contrast , EGRL has a simpler architecture using generic PG and EA to scale to large search spaces . For example , previous work with manual grouping operate at most in 5280 ≈ 10196 dimensional action space ( Mirhoseini et al. , 2017 ) , compared to ∼ 10358 for our BERT problem . Compared to pure PG based approaches , EGRL has significantly fewer hyperparameters to tune - primarily due to the reduced dependency on PG learning by using a population based search to handle sparse rewards . Another closely related work is REGAL ( Paliwal et al. , 2020 ) , which optimizes run-time and peakmemory via hardware placement . It also utilizes a graph representation with a genetic algorithm ( GA ) guided by RL . The RL agent predicts the parameters of GA - a form of indirect information transfer - while GA directly optimizes the final strategy . In contrast , our RL and EA components co-optimize the mapping strategies via direct information transfer ( policy migration ) and a shared replay buffer . REGAL assumes infinite bandwidth , whereas we train and validate entirely on physical hardware introducing specific mechanisms to incentivize compiler-valid mappings . This ensures that our solutions are performant under real-world operating conditions and closer to production-use . As a relatively new research field , we are challenged by the unavailability of reproducible code for prior work . The domain specific heuristics as described above render it difficult to apply algorithms designed for , say , chip placement to our memory mapping problem . Therefore , we adopt a stateof-the-art PG method as a baseline , since PG is a common central component of the above prior work . Classical Search Methods : Classical methods such as Simulated Annealing ( SA ) ( Kirkpatrick et al. , 1983 ) and genetic algorithms ( GA ) have also been studied for problems that have a similar combinatorial search complexity as memory placement . SA evolves new solution via small perturbations on existing solutions and retaining solutions that yield improved performance . A temperature parameter drives the exploration into new solution spaces . Our evolutionary algorithms ( EA ) ( Floreano et al. , 2008 ; Lüders et al. , 2017 ; Fogel , 2006 ; Spears et al. , 1993 ) improve on SA by systematically evolving new solutions within a population of solutions by performing mutations ( similar to SA ) and cross-over between pairs of solutions . Both methods are known to produce highly performant and stable solutions but are also significantly slower compared to Deep RL . In this work , we use EA both as a component of EGRL and also as a baseline . The PG components of EGRL produce fast , semi-performant solutions which then become anchors in the EA module . This essentially `` guides '' the EA to a performant solution by providing better anchors to search around . We demonstrate this via ablation studies that isolate the EA and PG components of EGRL . We also introduce Boltzmann chromosomes in the EA population - a set of stateless policies that directly perturb action proposals to accelerate exploration – with a temperature term that balances exploration and exploitation . This component is motivated by SA . Evolutionary RL : Our overall architecture builds on top of CERL ( Khadka and Tumer , 2018 ; Khadka et al. , 2019 ) which combines EA and PG . It diversifies exploration by allowing a population of EA policies to add data to a central replay buffer shared by a population of PG learners . We directly build on CERL because it has been shown to be effective in optimizing sparse feedback signals . Our memory mapping solution inherently relies on optimizing a very sparse feedback signal ( e.g. , latency ) that is obtained at the end of an inference run on a workload . For the PG portion of our architecture , we adopt Soft-Actor Critic ( SAC ) ( Haarnoja et al. , 2018 ) , the current state-of-the-art model-free algorithm developed for continuous high-dimensional settings . SAC uses an actor-critic architecture with separate networks for the policy and the Q-value function . A stochastic Gaussian policy enables it to use a maximum entropy objective ( Ziebart et al. , 2008 ) through which it demonstrates state-of-the-art results . We modify SAC to be compatible with our discrete action space .

Summary: The paper proposes Evolutionary Graph Reinforcement Learning to solve the memory placement problem. Main ideas are using GNN as the network architecture for reinforcement learning agents that look for more informed priors for evolutionary algorithms. Overall novelty of the paper comes from the neat combination of RL, EA, and GNN, and applying it to memory placement (ML for Systems).