  Larger deep learning models usually lead to higher model quality with an
ever-increasing GPU memory footprint. Although tensor checkpointing techniques
have been proposed to enable training under a restricted GPU memory budget, the
input tensor dynamics have been unexploited for optimizing performance while
reducing GPU memory footprint. Specifically, due to the diverse datasets and
subsequent data argumentation, the input tensor size per mini-batch is dynamic
during the training process, leading to a changing GPU memory footprint.
However, to leverage such input tensor dynamics in checkpointing, there are two
challenges to be solved. First, the checkpointing plan needs to be determined
during runtime due to the dynamics of input tensors. Second, the checkpointing
plan needs to be applied on the fly without significantly deteriorating the
performance.
  In this paper, we propose Mimose, an input-aware tensor checkpointing planner
respecting the memory budget while enabling efficient model training on GPU.
Mimose builds a lightweight but accurate prediction model of GPU memory usage
online, without pre-analyzing the model. It generates a tensor checkpointing
plan based on per-layer memory prediction and applies it to training progress
on the fly. It also adopts a caching strategy to avoid having to regenerate the
plan for repeated input size. Our experiments show that Mimose achieves
superior training throughput compared to state-of-the-art memory planners under
the same GPU memory budgets.
