# Rockmate: an Efficient, Fast, Automatic and Generic Tool for Re-materialization in PyTorch

arXiv: https://arxiv.org/abs/2307.01236 · HTML: https://mlsanity.com/paper/2307.01236/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 10 of its 23 references that are in the map.
- First posted on arXiv 3 Jul 2023.

## Builds on (most cited first, 10 of 10)

- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization](https://mlsanity.com/paper/1910.02653/) (2019)
- [Dynamic Tensor Rematerialization](https://mlsanity.com/paper/2006.09616/) (2020)
- [POET: Training Neural Networks on Tiny Devices with Integrated Rematerialization and Paging](https://mlsanity.com/paper/2207.07697/) (2022)
- [Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory](https://mlsanity.com/paper/1911.13214/) (2019)
- [TFLMS: Large Model Support in TensorFlow by Graph Rewriting](https://mlsanity.com/paper/1807.02037/) (2018)
- [A Graph Theoretic Framework of Recomputation Algorithms for Memory-Efficient Backpropagation](https://mlsanity.com/paper/1905.11722/) (2019)
- [Distributed Deep Learning Using Synchronous Stochastic Gradient Descent](https://mlsanity.com/paper/1602.06709/) (2016)
- [OLLA: Optimizing the Lifetime and Location of Arrays to Reduce the Memory Usage of Neural Networks](https://mlsanity.com/paper/2210.12924/) (2022)
- [Mimose: An Input-Aware Checkpointing Planner for Efficient Training on GPU](https://mlsanity.com/paper/2209.02478/) (2022)

## Built on by (most cited first, 4 of 4)

- [The Languini Kitchen: Enabling Language Modelling Research at Different Scales of Compute](https://mlsanity.com/paper/2309.11197/) (2023)
- [ProTrain: Efficient LLM Training via Memory-Aware Techniques](https://mlsanity.com/paper/2406.08334/) (2024)
- [Moonwalk: Inverse-Forward Differentiation](https://mlsanity.com/paper/2402.14212/) (2024)
- [TERRA: A Hierarchical Parallel Training and Memory Orchestration Framework for High-Resolution AI-based Earth Modeling](https://mlsanity.com/paper/2608.15211/) (2026)

## Abstract (arXiv)

We propose Rockmate to control the memory requirements when training PyTorch DNN models. Rockmate is an automatic tool that starts from the model code and generates an equivalent model, using a predefined amount of memory for activations, at the cost of a few re-computations. Rockmate automatically detects the structure of computational and data dependencies and rewrites the initial model as a sequence of complex blocks. We show that such a structure is widespread and can be found in many models in the literature (Transformer based models, ResNet, RegNets,...). This structure allows us to solve the problem in a fast and efficient way, using an adaptation of Checkmate (too slow on the whole model but general) at the level of individual blocks and an adaptation of Rotor (fast but limited to sequential models) at the level of the sequence itself. We show through experiments on many models that Rockmate is as fast as Rotor and as efficient as Checkmate, and that it allows in many cases to obtain a significantly lower memory consumption for activations (by a factor of 2 to 5) for a rather negligible overhead (of the order of 10% to 20%). Rockmate is open source and available at https://github.com/topal-team/rockmate.