"Broken" computation graph appears to use stale gradient

KFrank: <p>Hello Forum!</p>
<p>I tried a test where I expected to get “RuntimeError: Trying to backward<br>
through the graph a second time,” but instead I get results that make no<br>
sense to me.</p>
<p>The basic question:  What should happen if you move part of the building<br>
of the computation graph out of the optimization loop?</p>
<p>Here is the test script (with a couple of alternatives commented out):</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

a = torch.arange (9.0).reshape (3, 3)
a.requires_grad = True

opt = torch.optim.SGD ((a,), lr = 0.1)

b = a[:, 1].contiguous()   # a is optimized (with stale gradient) but b remains unchanged

# b = a[:, 1]              # works for first opt.step(), but then fails silently

for  i in range (3):
    # b = a[:, 1]          # works as expected inside the loop
    loss = (b[1] - 10)**2
    opt.zero_grad()
    loss.backward()
    opt.step()
    print ('a[1, 1]: ', a[1, 1])
    print ('b[1]:    ', b[1])
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">2.0.1
a[1, 1]:  tensor(5.2000, grad_fn=&lt;SelectBackward0&gt;)
b[1]:     tensor(4., grad_fn=&lt;SelectBackward0&gt;)
a[1, 1]:  tensor(6.4000, grad_fn=&lt;SelectBackward0&gt;)
b[1]:     tensor(4., grad_fn=&lt;SelectBackward0&gt;)
a[1, 1]:  tensor(7.6000, grad_fn=&lt;SelectBackward0&gt;)
b[1]:     tensor(4., grad_fn=&lt;SelectBackward0&gt;)
</code></pre>
<p>The first question is why, when <code>.backward()</code> is called the second time<br>
through the loop, does autograd not complain about part of the computation<br>
graph having been freed?  I would have thought that the first <code>.backward()</code><br>
would have freed the connection between <code>b</code> and <code>a</code> (which is not rebuilt inside<br>
of the optimization loop).</p>
<p>Why do we get the results we do?  It appears that <code>opt.step()</code> keeps using<br>
the gradient produced by the first <code>.backward()</code> and the subsequent calls<br>
don’t update it (but also don’t issue an error).</p>
<p>(The second – commented out – alternative without the <code>.contiguous()</code><br>
fails a little differently, but also without an error message.  In this case, <code>a</code><br>
is updated just once (and that change is reflected in <code>b</code>’s view into <code>a</code>) and<br>
then stays constant.)</p>
<p>I see this in versions 2.0.1 and 1.11.0, so it’s not an obvious one-off regression.</p>
<p>[Edit:  Adding another example script.]</p>
<p>I’ve added a modified version of the above script that highlights a particular<br>
facet of the issue.</p>
<p>After <code>loss</code> is first computed, it no longer changes when <code>a</code> changes.  This<br>
makes sense because <code>b</code>, the link between <code>a</code> and <code>loss</code> is not updated<br>
within the optimization loop; the line <code>b = a[:, 1].contiguous()</code> is only<br>
executed once before the loop begins.</p>
<p>Nonetheless, even though <code>loss</code> no longer depends on <code>a</code> inside of loop, a<br>
nonzero gradient flows back from <code>loss</code> to <code>a</code>.  (Note, constant loss, but a<br>
nonzero gradient, is clearly mathematically incorrect.)  It’s as if the part of<br>
the computation graph that connects <code>b</code> to <code>a</code> is still in operation, even though<br>
it should have been freed by the call to <code>loss.backward()</code>.</p>
<p>Here is a script that prints out <code>a.grad</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

a = torch.arange (9.0).reshape (3, 3)
a.requires_grad = True

opt = torch.optim.SGD ((a,), lr = 0.1)

b = a[:, 1].contiguous()   # a is optimized (with stale gradient) but b remains unchanged

for  i in range (2):
    loss = (b[1] - 10)**2
    print ('i:', i, ' loss =', loss)         # loss doesn't change (after first computed)
    print ('i:', i, 'before zero_grad()')
    print ('a = ...')
    print (a)
    print ('a.grad = ...')
    print (a.grad)
    opt.zero_grad()
    print ('i:', i, 'after zero_grad(), before backward()')
    print ('a = ...')
    print (a)
    print ('a.grad = ...')
    print (a.grad)
    loss.backward()
    print ('i:', i, 'after backward(), before step()')
    print ('a = ...')
    print (a)
    print ('a.grad = ...')
    print (a.grad)                           # grad is nonzero even though loss no longer depends on a
    opt.step()
    print ('i:', i, 'after step()')
    print ('a = ...')
    print (a)                                # a changes with every step
    print ('a.grad = ...')
    print (a.grad)
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">2.0.1
i: 0  loss = tensor(36., grad_fn=&lt;PowBackward0&gt;)
i: 0 before zero_grad()
a = ...
tensor([[0., 1., 2.],
        [3., 4., 5.],
        [6., 7., 8.]], requires_grad=True)
a.grad = ...
None
i: 0 after zero_grad(), before backward()
a = ...
tensor([[0., 1., 2.],
        [3., 4., 5.],
        [6., 7., 8.]], requires_grad=True)
a.grad = ...
None
i: 0 after backward(), before step()
a = ...
tensor([[0., 1., 2.],
        [3., 4., 5.],
        [6., 7., 8.]], requires_grad=True)
a.grad = ...
tensor([[  0.,   0.,   0.],
        [  0., -12.,   0.],
        [  0.,   0.,   0.]])
i: 0 after step()
a = ...
tensor([[0.0000, 1.0000, 2.0000],
        [3.0000, 5.2000, 5.0000],
        [6.0000, 7.0000, 8.0000]], requires_grad=True)
a.grad = ...
tensor([[  0.,   0.,   0.],
        [  0., -12.,   0.],
        [  0.,   0.,   0.]])
i: 1  loss = tensor(36., grad_fn=&lt;PowBackward0&gt;)
i: 1 before zero_grad()
a = ...
tensor([[0.0000, 1.0000, 2.0000],
        [3.0000, 5.2000, 5.0000],
        [6.0000, 7.0000, 8.0000]], requires_grad=True)
a.grad = ...
tensor([[  0.,   0.,   0.],
        [  0., -12.,   0.],
        [  0.,   0.,   0.]])
i: 1 after zero_grad(), before backward()
a = ...
tensor([[0.0000, 1.0000, 2.0000],
        [3.0000, 5.2000, 5.0000],
        [6.0000, 7.0000, 8.0000]], requires_grad=True)
a.grad = ...
None
i: 1 after backward(), before step()
a = ...
tensor([[0.0000, 1.0000, 2.0000],
        [3.0000, 5.2000, 5.0000],
        [6.0000, 7.0000, 8.0000]], requires_grad=True)
a.grad = ...
tensor([[  0.,   0.,   0.],
        [  0., -12.,   0.],
        [  0.,   0.,   0.]])
i: 1 after step()
a = ...
tensor([[0.0000, 1.0000, 2.0000],
        [3.0000, 6.4000, 5.0000],
        [6.0000, 7.0000, 8.0000]], requires_grad=True)
a.grad = ...
tensor([[  0.,   0.,   0.],
        [  0., -12.,   0.],
        [  0.,   0.,   0.]])
</code></pre>
<p>Context:  I got confused / tripped up by this when exploring what I thought<br>
would be a simple explanation / fix for the issue in the thread:</p>
<aside class="quote quote-modified" data-post="1" data-topic="186966">
  <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/winkawaks/48/63416_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/better-solution-of-backward-through-the-graph-a-second-time-problem/186966">Better solution of "backward through the graph a second time" problem?</a> <a class="badge-category__wrapper " href="/c/autograd/7"><span data-category-id="7" data-drop-close="true" class="badge-category " title="A category of posts relating to the autograd engine itself."><span class="badge-category__name">autograd</span></span></a>
  </div>
  <blockquote>
    Hi all, I have a specific need of my model. There is a learnable tensor – a, and I need some slices of it, each corresponding to a different loss function. The code is like this: 
import torch

a = torch.randn(3, 3)
a.requires_grad = True
b = a[:, :2].contiguous()
c = a[:, 1:].contiguous()
d = a[:, [0, 2]].contiguous()

loss_f = torch.nn.MSELoss()
opt = torch.optim.Adam([a], lr=0.1)
for i in range(100):
    opt.zero_grad()
    loss = loss_f(b, torch.randn_like(b)) + loss_f(c, torch.randn_like(c)…
  </blockquote>
</aside>

<p>Thanks for any insight!</p>
<p>K. Frank</p>

ptrblck: <p>I also stumbled across the issue and couldn’t explain it.<br>
CC <a class="mention" href="/u/alband">@albanD</a> for visibility, as I would expect to see an error instead of silent numerical errors.</p>

albanD: <blockquote>
<p>The first question is why, when <code>.backward()</code> is called the second time<br>
through the loop, does autograd not complain about part of the computation<br>
graph having been freed?</p>
</blockquote>
<p>This is because the part of the graph being shared is two ops: one indexing and one contiguous. neither of these saves any Tensor for backward. So nothing is there to be freed during the backward pass (and so we don’t detect the multiple backward calls).<br>
Note that what is being freed during the backward are only the saved Tensors, not the structure of the graph itself. So in practice you can backward as many times as you want through a given graph that doesn’t save any Tensor (but since what gets saved is not guaranteed, you shouldn’t rely on it).</p>
<blockquote>
<p>Why do we get the results we do?</p>
</blockquote>
<p>b doesn’t share memory with anything that is modified inplace and it is not recomputed so loss remains constant in the loop.<br>
Since the b doesn’t change, the loss doesn’t change and so you get the same gradient at every iteration, it is thus expected that a moves this way.</p>

KFrank: <p>Hi Alban!</p>
<aside class="quote no-group" data-username="albanD" data-post="3" data-topic="186980" data-full="true">
<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/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>neither of these saves any Tensor for backward. So nothing is there to be freed during the backward pass (and so we don’t detect the multiple backward calls).</p>
</blockquote>
</aside>
<p>I believe I understand your explanation of the behavior.</p>
<p>A couple of comments:</p>
<p>The example code I posted is, of course, logically incorrect from the<br>
perspective of backpropagation.</p>
<p>So one point of view is:  You write broken code, you get bad results.</p>
<p>However, in part because backpropagation errors can be easy to make, hard<br>
to debug, and sometimes hard to even notice, pytorch goes to considerable<br>
lengths to detect and warn about them.</p>
<p>This error is not exactly the same as a “backward through the graph a second<br>
time” error, but it’s at least a kissing cousin.  Would it be practical for pytorch<br>
to also detect this subclass of errors?</p>
<p>Thanks for the explanation!</p>
<p>K. Frank</p>

albanD: <p>Hey!</p>
<p>Indeed using <code>no_grad()</code> (which is what is done in optimizer.step() ) will hide things from autograd and might get you to compute surprising results.<br>
I don’t think we can make this an error as it would be BC-breaking pretty badly. In particular we don’t want to break existing code that does work today (when you do only a view outside of the main loop for example). But we can definitely add some detection (via a boolean on the Node) and warn_once if you backward again, even if no state if present. This will need to be a brand new mechanism though as the current one for errors will not work (it relies on states being unpacked, but you don’t have state here).</p>

KFrank: <p>Hi Alban!</p>
<p>Thanks for the follow up.</p>
<aside class="quote no-group" data-username="albanD" data-post="5" data-topic="186980" data-full="true">
<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/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>I don’t think we can make this an error as it would be BC-breaking pretty badly.</p>
</blockquote>
</aside>
<p>From where I sit, I don’t really see this as any kind of priority.  (Maybe it falls<br>
into the “might be nice” category.)</p>
<p>In fairness, in my example, I was purposely writing unusual “wrong” code.<br>
I haven’t ever tripped over this issue “in real life.”</p>
<p>Thanks again.</p>
<p>K. Frank</p>