Repository: pytorch
Issue: Dynamic tensor rematerialization
Labels: module: internals,module: checkpoint,feature,triaged

[Dynamic tensor rematerialization](https://arxiv.org/abs/2006.09616) (DTR) is a paper by Marisa Kirisame et al which is a new spin on checkpointing: like previous checkpointing systems, it reduces the amount of memory needed to run models by freeing and then recomputing activations on demand; unlike previous systems, it does so completely *online*, meaning it can be applied to existing models without modification and works well with PyTorch's eager model. The paper has gotten some attention: [MegEngine](https://github.com/MegEngine/MegEngine/wiki/Reduce-GPU-memory-usage-by-Dynamic-Tensor-Rematerialization) reimplemented the DTR algorithm for their framework and reproduced the benchmarks from the paper, although their conclusion was that memory fragmentation is still a big issue, and so as future work they want to introduce more static aspects (e.g., making use of a profile) to improve performance.

The [implementation on PyTorch](https://github.com/uwsampl/dtr-prototype) involves a number of pieces:

* An implementation of "checkpoint tensors": https://github.com/uwsampl/dtr-prototype/blob/eff53cc4804cc7d6246a6e5086861ce2b846f62b/dtr_code/dtr-implementation.patch#L762-L812 which manage tensors in the system and can be cleared out / recomputed, including a checkpoint dispatch key. The best description of the implementation aspects of this code are in Appendix E in https://arxiv.org/pdf/2006.09616.pdf
* An implementation, for every operator in PyTorch, of how to rematerialize the operation at some later point in time. This code is very regular and was written by hand by the paper authors. https://github.com/uwsampl/dtr-prototype/blob/eff53cc4804cc7d6246a6e5086861ce2b846f62b/dtr_code/dtr-implementation.patch#L1505
* New dispatcher functions for some backward functions (e.g., https://github.com/uwsampl/dtr-prototype/blob/eff53cc4804cc7d6246a6e5086861ce2b846f62b/dtr_code/dtr-implementation.patch#L3944 )
* Some generic PyTorch fixes
  * Some fixes for checkpointing unfriendly dispatch tables (e.g., https://github.com/uwsampl/dtr-prototype/blob/eff53cc4804cc7d6246a6e5086861ce2b846f62b/dtr_code/dtr-implementation.patch#L3494 )
  * Some miscellaneous fixes to replace uses of view with reshape ( https://github.com/uwsampl/dtr-prototype/blob/eff53cc4804cc7d6246a6e5086861ce2b846f62b/dtr_code/dtr-implementation.patch#L5117-L5139 ); is https://github.com/pytorch/pytorch/issues/55126 possibly related?

I'd like to see DTR live permanently as a PyTorch ecosystem library that is easy to install and toggle, similar to how AMP lived out of PyTorch core in APEX before eventually canonized into the main library. To do this, it is necessary to make some changes to how DTR is implemented (it is currently implemented as a patch on PyTorch):

1. Generic PyTorch fixes should be upstreamed into PyTorch (I believe the dispatch table adjustments and backward functions in native_functions.yaml should be noncontroversial; the need to replace view with reshape is probably some latent bug in the framework and worth investigating in more detail)
2. The rematerialization implementation should be rewritten in such a way that doesn't require every function to be written by hand. The simplest, albeit least efficient, way to do so is to implement it using a [boxed fallback](https://dev-discuss.pytorch.org/t/backend-fallbacks/195). Other alternatives include repurposing lazy tensor in some way (https://github.com/pytorch/rfcs/pull/18; a checkpoint tensor is a lazy tensor on its inputs and a possible evicted concrete tensor; NB lazy tensor is NOT YET landed to PyTorch master), or writing a code generation pass to autogenerate the code that was handwritten.
3. DTR needs to request a dispatch key for its tensor implementation; during prototyping a PrivateUse key can be used
4. At this point, the rest of DTR can entirely be implemented as a C++ extension https://pytorch.org/tutorials/advanced/cpp_extension.html

I also suspect that DTR can be implemented *entirely in Python* using the new dispatch to Python mechanism in https://github.com/pytorch/pytorch/issues/59049 The implementation here would proceed in this way:

1. Create a subclass of Tensor, CheckpointTensor, that holds both an underlying tensor (that may be evicted) and information on how to rematerialize the tensor
2. Define `__torch_dispatch__` on this tensor. This function will be called whenever an operation on a CheckpointTensor is run; the main algorithm of DTR can be implemented here.

Of course, this will be slower than the implementation of DTR from the paper. However, experiments with functorch with dispatch to Python show that the overhead is not too great, and I am guessing a sharp engineer could crank this out in a few days.

cc @ezyang @bhosmer @smessmer @ljk53 @bdhirsh