Dynamic Tensor Rematerialization

06/17/2020
by   Marisa Kirisame, et al.
102

Checkpointing enables training deep learning models under restricted memory budgets by freeing intermediate activations from memory and recomputing them on demand. Previous checkpointing techniques statically plan these recomputations offline and assume static computation graphs. We demonstrate that a simple online algorithm can achieve comparable performance by introducing Dynamic Tensor Rematerialization (DTR), a greedy online algorithm for checkpointing that is extensible and general, is parameterized by eviction policy, and supports dynamic models. We prove that DTR can train an N-layer linear feedforward network on an Ω(√(N)) memory budget with only 𝒪(N) tensor operations. DTR closely matches the performance of optimal static checkpointing in simulated experiments. We incorporate a DTR prototype into PyTorch just by interposing on tensor allocations and operator calls and collecting lightweight metadata on tensors.

READ FULL TEXT

page 2

page 6

page 14

page 25

page 26

research
09/06/2022

Mimose: An Input-Aware Checkpointing Planner for Efficient Training on GPU

Larger deep learning models usually lead to higher model quality with an...
research
03/30/2022

DELTA: Dynamically Optimizing GPU Memory beyond Tensor Recomputation

The further development of deep neural networks is hampered by the limit...
research
04/29/2021

Tuna: A Static Analysis Approach to Optimizing Deep Neural Networks

We introduce Tuna, a static analysis approach to optimizing deep neural ...
research
02/28/2022

DropIT: Dropping Intermediate Tensors for Memory-Efficient DNN Training

A standard hardware bottleneck when training deep neural networks is GPU...
research
10/09/2017

CTD: Fast, Accurate, and Interpretable Method for Static and Dynamic Tensor Decompositions

How can we find patterns and anomalies in a tensor, or multi-dimensional...
research
10/20/2018

The Ocean Tensor Package

Matrix and tensor operations form the basis of a wide range of fields an...
research
12/16/2021

A Static Analyzer for Detecting Tensor Shape Errors in Deep Neural Network Training Code

We present an automatic static analyzer PyTea that detects tensor-shape ...

Please sign up or login with your details

Forgot password? Click here to reset