Torch.utils.checkpoint does not support List inputs

username: hhxx
cooked: <p>I attempted to utilize <strong>torch.utils.checkpoint</strong> in my code, but I discovered that it does not support <strong>List</strong> inputs.</p>
<p>Below is the standard usage of <strong>torch.utils.checkpoint</strong>:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.utils.checkpoint as cp

class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.fc1 = nn.Linear(784, 1024)
        self.fc2 = nn.Linear(1024, 1024)
        self.fc3 = nn.Linear(1024, 512)
        self.fc4 = nn.Linear(512, 10)

    def forward(self, x):
        # Apply the first 2 layers
        x = nn.functional.relu(self.fc1(x))
        x = nn.functional.relu(self.fc2(x))
        # Apply the checkpointed layers
        x = cp.checkpoint(self._checkpointed_forward, x)

        # Apply the last 1 layer
        x = self.fc4(x)
        return x

    def _checkpointed_forward(self, x):
        x = nn.functional.relu(self.fc3(x))
        return x

# Create a random input tensor
x = torch.randn(2, 784)

# Create the neural network
model = MyNet()

# Compute the output
output = model(x)

# Compute the gradients
output.sum().backward()
# import pdb; pdb.set_trace()

# Print the gradients of the first layer
print(model.fc1.weight.grad)
</code></pre>
<p>If I modify the above code as follows (the main change is using a List in <strong>_checkpointed_forward</strong>):</p>
<pre><code class="lang-auto">class MyNet(nn.Module):
    def forward(self, x):
        ...
        x = cp.checkpoint(self._checkpointed_forward, [x])
        x = self.fc4(x)
        return x

    def _checkpointed_forward(self, x):
        x = nn.functional.relu(self.fc3(x[0]))
        return x
...
output.sum().backward()
print(model.fc1.weight.grad)
</code></pre>
<p>I observe that <strong>model.fc1.weight.grad = None</strong> in this case.<br>
Additionally, I set <strong>pdb.set_trace()</strong> in the <strong>class CheckpointFunction</strong>:</p>
<pre><code class="lang-auto">class CheckpointFunction(torch.autograd.Function):
    def backward(ctx, *args):
        import pdb; pdb.set_trace()
</code></pre>
<p>I notice that CheckpointFunction.backward is not even executed.</p>
<p>I am curious about why the code behaves in this manner.<br>
(torch.<strong>version</strong> = 1.10.0)</p>

username: hhxx
cooked: <p><a class="mention" href="/u/alband">@albanD</a> <a class="mention" href="/u/ptrblck">@ptrblck</a> <a class="mention" href="/u/richard">@richard</a> Could anyone give some advice?</p>

username: soulitzer
cooked: <aside class="quote no-group" data-username="hhxx" data-post="1" data-topic="193717">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/hhxx/48/21482_2.png" class="avatar"> hhxx:</div>
<blockquote>
<p>(torch.<strong>version</strong> = 1.10.0</p>
</blockquote>
</aside>
<p>List inputs are supported if you specify use_reentrant=False to activation checkpoint. (You may need a later version of PyTorch).<br>
The reason is that checkpoint with use_reentrant=True does not check that tensors hidden in lists are requires grad or not. At least one top-level tensor needs to have requires grad for use_reentrant=True, but this is not a constraint for use_reentrant=False. You can check the documentation here for more information <a href="https://pytorch.org/docs/stable/checkpoint.html#torch.utils.checkpoint.checkpoint" class="inline-onebox" rel="noopener nofollow ugc">torch.utils.checkpoint — PyTorch 2.1 documentation</a>.</p>