AutoVM: Accelerating Convolutional Neural Network Training with Actively Managed GPU Virtual Memory

The size of neural networks a GPU can train is limited by the GPU’s memory capacity. Although GPU virtual memory enables training arbitrarily large neural networks, such trainings are often accompanied by severe performance penalties. Furthermore, popular frameworks for constructing machine learning applications, like TensorFlow, have disabled using GPU virtual memory by default. We propose AutoVM, a software layer that can better manage GPU virtual memory in neural network training by incorporating our understandings of neural networks. AutoVM schedules data transfers between GPU and CPU memory to relieve the memory pressure on GPU; and in turn optimizes training speed. We have integrated AutoVM into TensorFlow such that existing machine learning applications can benefit from AutoVM with minimal effort. Our tests suggest that training VGG-19 using AutoVM can be at most 2.5× faster compared to using default Nvidia virtual memory.M.A.S