MONET: Modeling and Optimization of neural NEtwork Training from Edge to Data Centers
This work addresses the problem of efficient deep learning deployment for researchers and engineers by providing a tool for hardware-software co-design in training, though it appears incremental as it builds upon an existing inference framework.
The paper tackles the challenge of modeling neural network training on heterogeneous dataflow accelerators, which existing inference-focused tools fail to address, and introduces MONET, a framework that demonstrates capability in exploring design spaces and finding better hardware architectures for models like ResNet-18 and GPT-2.
While hardware-software co-design has significantly improved the efficiency of neural network inference, modeling the training phase remains a critical yet underexplored challenge. Training workloads impose distinct constraints, particularly regarding memory footprint and backpropagation complexity, which existing inference-focused tools fail to capture. This paper introduces MONET, a framework designed to model the training of neural networks on heterogeneous dataflow accelerators. MONET builds upon Stream, an experimentally verified framework that that models the inference of neural networks on heterogeneous dataflow accelerators with layer fusion. Using MONET, we explore the design space of ResNet-18 and a small GPT-2, demonstrating the framework's capability to model training workflows and find better hardware architectures. We then further examine problems that become more complex in neural network training due to the larger design space, such as determining the best layer-fusion configuration. Additionally, we use our framework to find interesting trade-offs in activation checkpointing, with the help of a genetic algorithm. Our findings highlight the importance of a holistic approach to hardware-software co-design for scalable and efficient deep learning deployment.