About the the implementation of activation checkpointing

syorami: <p>I read the source code and found that when using activation checkpointing, the block forward is in <code>torch.no_grad</code> and the outputs are supposed to have <code>requires_grad</code> set to False. However, after I called the checkpoint, the output tensor required gradient, and I couldn’t figure out why this happened. Hope anyone could help me on this.</p>

ptrblck: <p>I’m not sure what your exact use case is, but the <a href="https://pytorch.org/docs/stable/checkpoint.html#torch.utils.checkpoint.checkpoint_sequential">docs</a> for <code>checkpoint_sequential</code> mention:</p>
<blockquote>
<p>All segments except the last will run in <a href="https://pytorch.org/docs/stable/generated/torch.no_grad.html#torch.no_grad"><code>torch.no_grad()</code></a> manner, i.e., not storing the intermediate activations. The inputs of each checkpointed segment will be saved for re-running the segment in the backward pass.</p>
</blockquote>
<p>Could this explain the behavior you are seeing?</p>

syorami: <p>hi <a class="mention" href="/u/ptrblck">@ptrblck</a>  I’m not facing any problem on its usage but just curious about the what is actually going on. Here is a small demo to reproduce the results:</p>
<pre><code class="lang-auto">import torch
from torch import nn
from torch.utils.checkpoint import checkpoint

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(100, 100)
        self.blocks = nn.Sequential(*[
            nn.Linear(100, 100),
            nn.Linear(100, 100)
        ])

    def forward(self, x):
        x = self.linear(x)
        x1 = x2 = x

        with torch.no_grad():
            for block in self.blocks:
                x1 = block(x1)
                print(f'no_grad requires_grad: {x1.requires_grad}')

        for block in self.blocks:
            x2 = checkpoint(block, x2)
            print(f'checkpoint requires_grad: {x2.requires_grad}')

        return x2

model = Model()
x = torch.randn((2, 100))
outs = model(x)
</code></pre>
<p>The outputs are as follows:</p>
<pre><code class="lang-auto">no_grad requires_grad: False
no_grad requires_grad: False
checkpoint requires_grad: True
checkpoint requires_grad: True
</code></pre>
<p>For checkpoint outputs, the <code>requires_grad</code> are set to True and this differs from the description of activation checkpoint.</p>

ptrblck: <p>Thanks for the code snippet as I see where the unexpected behavior is coming from.<br>
If you add a debug print statement into the <code>forward</code> of <code>CheckpointFunction</code> <a href="https://github.com/pytorch/pytorch/blob/0cc0e5ef65c09983e6e7127de9d16176a1219e91/torch/utils/checkpoint.py#L107">here</a>, you’ll see that <code>output</code> indeed does not require gradients anymore. However, the returned tensor returns <code>True</code> in its <code>requires_grad</code> attribute again as it seems to be a properly of an <code>autograd.Function</code>:</p>
<pre><code class="lang-python">class MyCheckpointFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        print(x.requires_grad)
        
        with torch.no_grad():
            y = x + 1
        print(y.requires_grad)
        return y
    
f = MyCheckpointFunction.apply
x = torch.randn(1, requires_grad=True)
out = f(x)
print(out.requires_grad)
</code></pre>
<p>Output:</p>
<pre><code class="lang-python">True
False
True
</code></pre>

syorami: <p>Yes! It is exactly the same misbehaviour I’ve noticed in previous experiments but I couldn’t find where the attributes are modified. Do you have any idea about this?</p>
<p>BTW, I found that the <code>requires_grad</code> attribute of the output tensor also depends on the inputs. Once I set the <code>requires_grad=False</code> for inputs, the outputs tensors are also <code>requires_grad=False</code>.</p>

ptrblck: <p>I don’t think this is a misbehavior, but exactly how custom <code>autograd.Function</code>s are designed and have to act on input tensors to avoid breaking the computation graph.<br>
E.g. a custom <code>autograd.Function</code> can be used to implement any operation from 3rd party libs (such as numpy) which Autograd cannot track. If the output tensor would not require gradients anymore, your layer would never work inside a proper computation graph.</p>

syorami: <p>Thanks for your explanation. This makes things much clear for me.</p>