Dynamic Tensor Rematerialization(DTR)

• Save memory for NN by dynamically discarding and recomputing intermediate results at runtime.
• By being smart about what to keep and what to discard, train larger models under a tight budget.
• Save 3x memory for 20% extra compute time - train models 3x as large!
• DTR Makes no assumptions about the model, and does not require a static graph.
• DTR can be easily implemented into different deep learning frameworks.
• DTR still achieves amazing space-time tradeoff!

Adoption:

Inspired:

Technical Dive

0. Gradient Checkpointing is a technique to save memory for Deep Neural Network training.
1. Or more generally, for reverse-mode automatic differentiation.
2. However, memory planning is np-complete.
3. Checkpointing also has to deal with program with arbitary control flow.
4. To combat this, previous work made different restriction which sacrifice performance or usability.
5. Some works models the program as a stack machine with no heap...
6. And suffers performance degradation when the assumption is broken!
7. (For example, NN with highway connection/branching).
8. Other works use an ILP solver, which consume lots of time to find the optimal memory planning.
9. And can only be used for program/framework without control flow, posing problem for real world adoption.
10. Additionally, gradient checkpointing couple derivative calculation, with memory saving by recomputing.
11. This add complexity and limit applications range.
12. DTR tackles the problems above by planning the memory greedily at runtime, instead of as a compiler pass.
13. Essentially, DTR always pair each tensor with a thunk(think lazy evaluation), so all tensors can be recomputed after evicting.
14. This solves the control flow and stack machine issue, as we do not model the program in anyway!
15. However, with a novel cache eviction policy, we are still able to achieve great performance.
16. Suprisingly DTR 'accidentally' discovered the O(n) time O(n^0.5) space, and the O(n log n) time O(log n) space scheme.
17. These two are both well-known result.
18. Furthermore DTR switch smoothly between normal execution without recomputation, O(n^0.5) space scheme, and O(log n) space scheme
.