In our model trainings, we use Activation Checkpointing [1]} which is the technique of trading compute for memory. It works by discarding all model activations during the forward pass except the layers that have been configured to be “checkpointed”. During the backward pass (backpropagation), the forward pass on a part of the model (between two checkpointing layers) is re-computed. While this technique can help increase the batch size (leading to more compute to be overlapped with communication which leads to more efficient training), one downside is that manual configuration / tuning of which layers should be “checkpointed” is needed. This can be time-consuming and often hard to find the optimal checkpointing state. Further, for the models that are hard to fit in memory, it can become very difficult to perform manual tuning.
