Tensor.requires_grad in torch.autograd.Function

hhxx: <p>I attempted to implement a custom function as outlined in the <a href="https://pytorch.org/docs/stable/notes/extending.html#extending-torch-autograd" rel="noopener nofollow ugc">tutorial</a>:</p>
<pre><code class="lang-auto"># Inherit from Function
class LinearFunction(Function):

    # Note that forward, setup_context, and backward are @staticmethods
    @staticmethod
    def forward(input, weight, bias):
        output = input.mm(weight.t())
        print('inputs.requires_grad: ', input.requires_grad)
        print('weight.requiers_grad: ', weight.requires_grad)
        print('output.requires_grad: ', output.requires_grad)
        if bias is not None:
            output += bias.unsqueeze(0).expand_as(output)
        return output

    @staticmethod
    # inputs is a Tuple of all of the inputs passed to forward.
    # output is the output of the forward().
    def setup_context(ctx, inputs, output):
        input, weight, bias = inputs
        ctx.save_for_backward(input, weight, bias)

    # This function has only a single output, so it gets only one gradient
    @staticmethod
    def backward(ctx, grad_output):
        # This is a pattern that is very convenient - at the top of backward
        # unpack saved_tensors and initialize all gradients w.r.t. inputs to
        # None. Thanks to the fact that additional trailing Nones are
        # ignored, the return statement is simple even when the function has
        # optional inputs.
        input, weight, bias = ctx.saved_tensors
        grad_input = grad_weight = grad_bias = None

        # These needs_input_grad checks are optional and there only to
        # improve efficiency. If you want to make your code simpler, you can
        # skip them. Returning gradients for inputs that don't require it is
        # not an error.
        if ctx.needs_input_grad[0]:
            grad_input = grad_output.mm(weight)
        if ctx.needs_input_grad[1]:
            grad_weight = grad_output.t().mm(input)
        if bias is not None and ctx.needs_input_grad[2]:
            grad_bias = grad_output.sum(0)

        return grad_input, grad_weight, grad_bias

x = torch.rand(2, 128, requires_grad=True)
weight = torch.rand(256, 128, requires_grad=True)
bias = troch.rand(256, requires_grad=True)
out = LinearFunction.apply(x, weight, bias)
print('out.requiers_grad: ', out.requiers_grad)
</code></pre>
<p>Upon executing the above code, the following outputs were obtained:</p>
<pre><code class="lang-auto">inputs.requires_grad: True
weight.requires_grad: True
output.requires_grad: False
out.requires_grad: True
</code></pre>
<p>The issue arises from the operation <strong>output = input.mm(weight.t())</strong>, generating a tensor without gradients despite both input and weight tensors having gradients enabled.</p>
<p>I assume that the absence of gradients in the output might be due to the forward function being invoked like:</p>
<pre><code class="lang-auto">def some_func():
    with torch.no_grad:
        forward(...)
</code></pre>
<p>If my assumption holds, I have three additional questions:</p>
<ol>
<li>
<p>Where can I find the definition of the aforementioned <strong>some_func()</strong>?</p>
</li>
<li>
<p>Why the <strong>output.requires_grad</strong> = False in the forward function, but <strong>out.requires_grad=True</strong> outside? What additional operation is applied after executing the forward() function?</p>
</li>
<li>
<p>What is the purpose of such a design? Is it intended to conserve GPU memory?</p>
</li>
<li>
<p>As far as I know, there is a technique called <strong>Checkpoint</strong> employed to optimize memory usage, which is detailed in <a href="https://pytorch.org/docs/stable/checkpoint.html" rel="noopener nofollow ugc">torch.utils.checkpoint</a>. It involves the omission of gradients during the forward process, with a subsequent re-computation during the backward process. I am curious about the distinctions and connections between <code>torch.autograd.Function</code> and <code>torch.utils.checkpoint</code>, particularly regarding their treatment of gradients in the forward process.</p>
</li>
</ol>

hhxx: <p>Additionally, I have reviewed the implementation of <a href="https://github.com/pytorch/pytorch/blob/0978482afa69118da1a986a4edec3acea01d2c6d/torch/utils/checkpoint.py#L223" rel="noopener nofollow ugc">torch.utils.checkpoint</a>:</p>
<pre><code class="lang-auto">class CheckpointFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, run_function, preserve_rng_state, *args):
        ...
    def backward(ctx, *args):
        ...
</code></pre>
<p>This class, <strong>CheckpointFunction</strong>, is a subclass of <strong>torch.autograd.Function</strong>. I have two additional questions regarding the <strong>checkpoint</strong>:</p>
<ol start="5">
<li>
<p>What techniques does the <strong>CheckpointFunction</strong> class employ to conserve GPU memory? Is it primarily achieved by omitting gradients in the forward process (though the original <strong>torch.autograd.Function</strong> also seems to run in a <code>no_grad()</code> environment)?</p>
</li>
<li>
<p>In the forward function of the <strong>CheckpointFunction</strong> class, there are the <a href="https://github.com/pytorch/pytorch/blob/0978482afa69118da1a986a4edec3acea01d2c6d/torch/utils/checkpoint.py#L261" rel="noopener nofollow ugc">following codes</a>:</p>
</li>
</ol>
<pre><code class="lang-auto">    with torch.no_grad():
        outputs = run_function(*args)
    return outputs
</code></pre>
<p>Is it redundant to define a new <code>no_grad()</code> environment? In reference to my first question, it appears that the forward function in the <strong>torch.autograd.Function</strong> class has already been wrapped in the <code>torch.no_grad()</code> statement.</p>

albanD: <p>Hey!</p>
<ol>
<li>The implementation is quite tricky and split between python and C++ I’m afraid.</li>
<li>This is expected behavior: you are providing the backward formula for the ops happening here (via the backward() method) so we don’t need to track autograd information during the forward. When the output is being returned, we will attach the appropriate autograd metadata to it so that it requires_grad and it will call your custom backward().</li>
<li>This is done mainly for performance: no need to create the autograd graph during the forward since we’re never going to use it. There are other more subtle implications like not saving things for backward twice etc that are also important.</li>
<li>checkpoint will omit intermediary buffers (save_for_backward), not gradients. And they are very much independent. As you can see in the utils.checkpoint code, we are transitioning to not use autograd.Function to implement AC (Activation Checkpointing) anymore.</li>
<li>It is by only saving for backward the input.</li>
<li>It is indeed redundant.</li>
</ol>

hhxx: <p>Hi <a class="mention" href="/u/alband">@albanD</a>, thanks for your detailed response. Most of my questions have been resolved.</p>
<p>I have one remaining question regarding your second answer: <strong>We will attach the appropriate autograd metadata to it so that it requires_grad</strong></p>
<p>My question is: Is the function responsible for attaching the grad metadata also implemented in the C++ source code, like the aforementioned <strong>some_func()</strong>?</p>
<p>Looking forward to your reply <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

albanD: <p>Yes it is in c++<br>
It is here exactly if you’re curious: <a href="https://github.com/pytorch/pytorch/blob/d9c0e37bab9462c18508a594659cd34a66abfe1e/torch/csrc/autograd/custom_function.cpp#L355-L362">https://github.com/pytorch/pytorch/blob/d9c0e37bab9462c18508a594659cd34a66abfe1e/torch/csrc/autograd/custom_function.cpp#L355-L362</a> <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>