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

By Benoit Steiner, Mostafa Elhoushi, Jacob Kahn, James Hegarty. 2022

11 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Hardware and Architecture, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

1. Christopher R'e
2. Ce Zhang
3. Ion Stoica
4. Joseph E. Gonzalez
5. Percy Liang
6. Daniel Y. Fu
7. Beidi Chen
8. Lianmin Zheng

In Computer Vision and Pattern Recognition

1. Christopher R'e
2. Joseph E. Gonzalez
3. Percy Liang
4. Ion Stoica
5. Daniel Y. Fu
6. Ce Zhang
7. Trevor Darrell
8. Luc Van Gool

In Hardware and Architecture

1. Daniel Y. Fu
2. Percy Liang
3. Ce Zhang
4. Christopher R'e
5. Ion Stoica
6. Beidi Chen
7. Joseph E. Gonzalez

In Efficiency & Compression

1. Ion Stoica
2. Daniel Y. Fu
3. Joseph E. Gonzalez
4. Percy Liang
5. Ce Zhang
6. Clark W. Barrett
7. Christopher R'e
8. Lianmin Zheng

Citations in the corpus

• Efficient Memory Management for Large Language Model Serving with PagedAttention (2023, background): “OLLA [48] optimizes the lifetime and location of tensors to reduce fragmentation, but it does not do finegrained block-level management or online serving.”
• High-throughput Generative Inference of Large Language Models with a Single GPU (2023, background): “Within broader domains, memory optimizations and offloading have been studied for training (Huang et al., 2020; Ren et al., 2021; Steiner et al., 2022) and linear algebra (Jia-Wei & Kung, 1981; Demmel, 2013).”