Gradient checkpointing not saving memory

username: Julian_Buchel
cooked: <p>Hi,</p>
<p>I implemented a triton kernel that is called inside a <code>torch.autograd.Function</code>. When using it in my training, I got an OOM. This led me to believe that activation checkpointing doesn’t work with <code>torch.autograd.Function</code>. I guess when I explicitly save tensors to the <code>ctx</code> they actually get saved and <code>torch.utils.checkpoint</code> doesn’t do anything about that. Is that correct? If so, how should I adapt my code to also support checkpointing?</p>
<p>Also, I played around with the <code>checkpoint</code> function using this script on one V100 GPU (32GB of VRAM).</p>
<pre><code class="lang-auto">import torch
from torch.nn import Linear
from torch.utils.checkpoint import checkpoint
from torch.autograd import Function

class LinearFun(Function):
    @staticmethod
    def forward(ctx, inp, weight):
        ctx.save_for_backward(inp, weight)
        return inp @ weight.T
    @staticmethod
    def backward(ctx, grad_output):
        inp, weight = ctx.saved_tensors

        grad_w = grad_output.T @ inp
        grad_inp = grad_output @ weight
        return grad_inp, grad_w

class NewLinear(Linear):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def forward(self, inp):
        return LinearFun.apply(inp, self.weight)

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.l1 = Linear(1024, 1024)
        self.l2 = Linear(1024, 1024)
        self.l3 = Linear(1024, 1)
        # self.l1 = NewLinear(1024, 1024)
        # self.l2 = NewLinear(1024, 1024)
        # self.l3 = NewLinear(1024, 1)
    
    def forward(self, inp):
        inp = self.l1(inp)
        inp = self.l2(inp)
        inp = self.l3(inp)
        return inp


model = Model().cuda()
inp = torch.randn(int(1e6), 1024, device="cuda")

# &lt;---- USE CHECKPOINTING
out = checkpoint(model, inp, use_reentrant=False)
# &lt;---- NOT USING CHECKPOINTING
# out = model(inp)

loss = out.sum()
print('========= before backward =========')
loss.backward()
print('========= after backward =========')
</code></pre>
<p>I tweaked the number of inputs to <code>int(1e6)</code> since that brings both scenarios to an OOM. But I don’t see that the checkpointing allows me to got to higher input dimensions.<br>
Am I missing something?</p>

username: soulitzer
cooked: <aside class="quote no-group" data-username="Julian_Buchel" data-post="1" data-topic="205575">
<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/julian_buchel/48/37532_2.png" class="avatar"> Julian_Buchel:</div>
<blockquote>
<p>I guess when I explicitly save tensors to the <code>ctx</code> they actually get saved and <code>torch.utils.checkpoint</code> doesn’t do anything about that.</p>
</blockquote>
</aside>
<p>Your autograd custom Function code looks fine. Saving to ctx via <code>ctx.save_for_backward(*tensors)</code> should be what you want.</p>
<p>However, when applying checkpointing, you do not want to apply it to the entire model, because during backward all the tensors intended to be saved for backward in that region would end up being materialized at once anyway.</p>
<p>To avoid the peak memory spike that you are still observing, you’ll want to apply checkpoint piece by piece so that during backward rematerialization also occurs piece by piece.</p>

username: Julian_Buchel
cooked: <p>Ok, so then in this example, there is actually no benefit you get from activation checkpointing? Or how would you re-write this so that the peak memory consumption is reduced?</p>

username: Julian_Buchel
cooked: <p>Also, I was wondering whether torch will avoid saving the tensors to the context of the function when gradient checkpointing is enabled.</p>

username: soulitzer
cooked: <p>Yes, checkpoint can interpose logic to avoid saving when you use the ctx.save_for_backward API</p>

username: Julian_Buchel
cooked: <p>Ok, nice. Somewhat unrelated, but when calling <code>ctx.save_for_backward</code>, it doesn’t save tensors when I am in <code>eval</code> mode, right? Or do I have to take care of that explicitly?</p>

username: soulitzer
cooked: <p>Calling mod.eval() doesn’t actually disable recording of the autograd graph/saving for backward. You might be thinking of torch.no_grad, or mod.requires_grad_(False).<br>
This might be of interest to you: <a href="https://pytorch.org/docs/stable/notes/autograd.html#locally-disabling-gradient-computation" class="inline-onebox" rel="noopener nofollow ugc">Autograd mechanics — PyTorch 2.3 documentation</a></p>

username: soulitzer
cooked: <p>It’s tricky because there are few layers in this example, but maybe grouping the first linears into a single checkpoint call, and not checkpointing the last linear would be helpful.</p>

username: Julian_Buchel
cooked: <p>Hi. Yes I forgot that we record the inputs in eval mode. In the meanwhile I solved my bug. It seems that everything is handled well when an autograd.Function is used. Thanks <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>