Note it’s been a while since I last looked at this. Also, I think the source for torch.utils.CheckpointFunction might be slightly different now, though hopefully not in a way that affects this question.


Say we didn’t detach the input, when we recompute the output and run the backward, the backward propagation would propagate past the input, as it is still attached to the gradient graph from the prior computations that created it. However, the function we are implementing (CheckpointFunction's backward) is supposed to be just one operator in that gradient graph, and return just its backward, so we do not want to cause that side effect.
With respect to whether detaching duplicates the underlying buffers, I don’t remember what I found but it should be simple to verify by experimentation (the tensor objects have a pointer to the underlying buffer; make a detached copy and check them for equality).


This was actually really simply and took far too long for me to figure out. The API of Autograd is assuming that the backward function returns one gradient for each argument given to the forward function, excluding the first context argument. If we look at the corresponding forward of CheckpointFunction, we see there are three inputs, the first two being ‘pseudo inputs’ that configure the checkpointing rather than being actual tensors, but we still need to account for them when returning the gradients from the backward function, hence (None, None) + grads.


I’m afraid I can’t recall how I figured this out. Sorry.


I think my analysis here was correct; the tensors to be checkpointed only on the first recomputation were being checkpointing immediately.
To get around this, I added a recomputation_depth parameter to the forward function of CheckpointFunction, and only checkpointed the input if the depth was 0. The forward of the Drop module carries its own recomputation_depth parameter and, in its forward, decrements it before invoking the checkpoint function.
Note that, as explained in 2., we now must return (None, None, None) + grads in the backward, due to the additional argument given to the forward.
This successfully allows you to encode multiple recomputations using these Drop modules, for example by doing:
torch.nn.Sequential(
    layer1,
    Drop(torch.nn.Sequential(
        layer2,
        Drop(torch.nn.Sequential(
            layer3,
            layer4,
        2)),          # <----- recursion depth of inner drop
        layer5,
    1                 # <----- recursion depth of outer drop
)

where layer 1 is not recomputed; 2 and 5 are recomputed once; and 3 and 4 are recomputed twice.


Hope this helps.