Saved tensor hooks checkpoint implementation cannot robustly clear storage

[rohan-varma]

🐛 Describe the bug

In the saved tensor hooks based checkpointing approach (

pytorch/torch/utils/checkpoint.py

Line 349 in 386b398

storage: Dict[int, Optional[torch.Tensor]] = {}

), when autograd needs to unpack an activation, it potentially re-runs the forward to recompute all activations, and then returns the activation for the index it is unpacking.

However, we currently do a storage.pop() for this to ensure we don't hold references to the tensor after the backward is over. This raises the issue that if the same tensor is unpacked twice, without a pack in between, we'll run into an error. A (silly) example repro is here:

pytorch/test/test_autograd.py

Line 4598 in 386b398

def test_checkpointing_without_reentrant_custom_function_raises(self):

Another concern is if a tensor is packed by autograd but never unpacked, thus leading to storage leaking. Although, we are unsure if this can occur in practice.

Versions

main

cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @lezcano @Varal7

[rohan-varma]

cc @albanD @ezyang