  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.
