Offload models to CPU using autograd.Function

cliven: <p>I was wondering if it was possible to do something like the following, where I try to load the model from CPU -&gt; GPU before the computation and send it back after:</p>
<pre><code class="lang-python">import torch
from torch import nn

DUMMY = torch.empty((), requires_grad=True)
class Clive(torch.autograd.Function):
    @staticmethod
    def forward(ctx, layer, dummy, inputs):
        layer.to(0)
        inputs = inputs
        with torch.enable_grad():
            outputs = layer(inputs)
        ctx.layer = layer
        ctx.inputs = inputs
        ctx.outputs = outputs
        layer.to('cpu')
        return outputs.detach()

    @staticmethod
    def backward(ctx, *grad_outputs):
        ctx.layer.to(0)

        diff_with = ctx.layer.parameters()
        if ctx.inputs.requires_grad:
            diff_with = itertools.chain([ctx.inputs], diff_with)

        grads = torch.autograd.grad(outputs=(ctx.outputs,), inputs=diff_with, grad_outputs=grad_outputs)
        
        if ctx.inputs.requires_grad:
            input_grad, *grads = grads
        else:
            input_grad = None

        for g, l in zip(grads, ctx.layer.parameters()):
            l.backward(g)
        ctx.layer.to('cpu')
        return (None, None, input_grad)

m1 = nn.Linear(10, 10)
x = torch.randn(10, 10, device=0)
Clive.apply(m1, DUMMY, x).sum().backward()
print(m1.weight.grad)
</code></pre>
<p>Unfortunately this errors with <code>One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.</code>. I’m still working to understand this error, can anyone dumb it down for me? If I get rid of the <code>.to(...)</code> method calls it works ok.</p>
<p>(I don’t really know if this is practical/useful, I’m just curious if I can actually do it :D)</p>

albanD: <p>Hi,</p>
<p>Moving across devices is differentiable. So you don’t have to use custom Function if you just want to use different devices.</p>

cliven: <p>I was also interested in offloading the computed outputs to the CPU (like gradient checkpointing but not recomputing), but you’re right I can just wrap the model.</p>
<p>Speaking of, is there something obvious I am doing wrong here? It seems to fail if I try to run another backward pass when there’s already gradients allocated.</p>
<pre><code class="lang-auto">import torch
from torch import nn
class Wrapped(nn.Module):
    def __init__(self):
        super().__init__()
        self.l = nn.Linear(10, 10)
        
    def forward(self, x):
        self.l.to(0)
        ans = self.l(x)
        self.l.to('cpu')
        return ans

inp = torch.randn(10, 10, device=0)
w = Wrapped()

# Standard forward/backward
w(inp).sum().backward()

# Clear grads and do it again
for p in w.parameters():
    p.grad = None
w(inp).sum().backward()

# Zero grads like we should do, fails with "Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!"
w.zero_grad()
w(inp).sum().backward()
</code></pre>

cliven: <p>Actually it doesn’t seem to work for my use case - I have to reload the model to GPU before running the backward pass. I’ve mimicked a trivial training loop here:</p>
<pre><code class="lang-auto">import torch
from torch import nn

m = nn.Linear(10, 10)
inputs = [
    torch.randn(10, 10, device=0),
    torch.randn(10, 10, device=0)
]
for input in inputs:
    m.to(0)
    activation = m(input)
    m.to('cpu')
    loss = activation.sum()
    # m.to(0) &lt;--- without this copy, it doesn't work
    loss.backward()    
    # m.to('cpu')
    m.zero_grad()
</code></pre>
<p>Obviously in this example this is fine, but what I’m aiming to do is only load the layer’s weight into GPU memory when they are needed (e.g. forward and backward). That was what I was going for with the autograd Function.</p>

albanD: <p>Hi,</p>
<p>You shouldn’t use the <code>Module.to()</code> functions inside your forward. It won’t work well with AD.<br>
Only the <code>Tensor.to()</code> that you were mentioning above will.</p>
<p>The reason for this is that the Module Parameters are always leaf Tensor. And so cannot be updated in a differentiable manner.</p>
<p>Note that more generally, I don’t think your approach will help wrt to GPU memory as the autograd needs to keep values to be able to compute the backward pass. And your model’s weights are part of them. So even if you move the model to the cpu, these buffers will stay on the GPU.</p>

ralzq01: <p>Hi, <a class="mention" href="/u/alband">@albanD</a></p>
<p>Thanks for the explanation! I’m curious that is there any simple way that I can offload the weight / intermediate output to CPU memory while getting back to GPU memory when it needs to (e.g., during backward, or at other place some op needs it)? Here I mean the true offloading that can save GPU memory.</p>
<p>I’m trying to make it by autograd function, like in-placement update on tensor data, but still not work (some backward error on gradient format)</p>

albanD: <p>I’m afraid there is no simple way to do this today.<br>
We are working on adding this though: <a href="https://github.com/pytorch/pytorch/issues/58659" class="inline-onebox">SavedVariable default hooks · Issue #58659 · pytorch/pytorch · GitHub</a> (this should be done in 1-2 months).</p>
<p>The workaround with using custom Function would work but you will need to save on the ctx foo.cpu() and not use save_for_backward. And then in the backward formula do ctx.foo.cuda().<br>
But that means that you have to rewrite most of the backward formulas which is very painful.</p>

ralzq01: <p>Great thanks! Looking forward to the new functionality!</p>