OLLA: Optimizing the Lifetime and Location of Arrays to Reduce the Memory Usage of Neural Networks

TL;DR

OLLA is an algorithm that optimizes tensor lifetime and placement during neural network training, significantly reducing memory usage without altering models or training procedures, enabling faster and more memory-efficient training.

Contribution

OLLA introduces a novel ILP-based approach to optimize tensor memory management in neural networks without model modifications.

Findings

Reduces neural network memory usage by one-third on average.

Operates efficiently, taking minutes or seconds for large networks.

Does not require changes to existing training procedures.

Abstract

The size of deep neural networks has grown exponentially in recent years. Unfortunately, hardware devices have not kept pace with the rapidly increasing memory requirements. To cope with this, researchers have turned to techniques such as spilling and recomputation, which increase training time, or reduced precision and model pruning, which can affect model accuracy. We present OLLA, an algorithm that optimizes the lifetime and memory location of the tensors used to train neural networks. Our method reduces the memory usage of existing neural networks, without needing any modification to the models or their training procedures. We formulate the problem as a joint integer linear program (ILP). We present several techniques to simplify the encoding of the problem, and enable our approach to scale to the size of state-of-the-art neural networks using an off-the-shelf ILP solver. We experimentally demonstrate that OLLA only takes minutes if not seconds to allow the training of neural networks using one-third less memory on average.