Year: 2019
Type: preprint
Abstract: We formalize the problem of trading-off DNN training time and memory requirements as the tensor rematerialization optimization problem, a generalization of prior checkpointing strategies. We introduce Checkmate, a system that solves for optimal rematerialization schedules in reasonable times (under an hour) using off-the-shelf MILP solvers or near-optimal schedules with an approximation algorithm, then uses these schedules to accelerate millions of training iterations. Our method scales to complex, realistic architectures and is hardware-aware through the use of accelerator-specific, profile-based cost models. In addition to reducing training cost, Checkmate enables real-world networks to be trained with up to 5.1x larger input sizes. Checkmate is an open-source project, available at https://github.com/parasj/checkmate.
Source: arXiv (Cornell University)
Authors: Jain, Paras, Jain, Ajay, Nrusimha, Aniruddha, Gholami, Amir, Abbeel, Pieter, Keutzer, Kurt, Stoica, Ion, Gonzalez, Joseph E.
Institution: University of California, Berkeley
Language: English
Subfield: Hardware and Architecture
Field: Computer Science
Domain: Physical Sciences
Open Access status: green
OpenAlex ID: W2979044977