RuntimeError: Error while setting up backward hooks. Please open an issue with a code sample to reproduce this

AlphaBetaGamma96: <p>Hi All,</p>
<p>I’ve been experimenting with creating a custom optimizer class. This optimizer requires extra information in order to work, which requires me to calculate more than 1 loss. I can talk a little bit more about this optimizer but in short, it’s a modified version of KFAC.</p>
<p>From using KFAC, you precondition your gradients via using information gained from <code>register_forward_pre_hook</code> and <code>regsiter_full_backward_hook</code>. However, the <code>grad_output</code> from <code>register_full_backward_hook</code> I require isn’t the same as what’s returned from the loss function, so I get round this issue by using 2 loss functions and caching the needed values. It goes something like this,</p>
<ul>
<li>
<p>add <code>register_foward_pre_hook</code> and <code>register_full_backward_hook</code> to my model for all layers.</p>
</li>
<li>
<p>call my loss function, and cache the correct gradient I need for my model. Although, the backward hooks have the wrong grad_output values. (The forward hook is fine as its independent of the loss function at hand).</p>
</li>
<li>
<p>call an auxiliary (or secondary) loss function (which returns the right grad_output for the backward hooks but has the wrong gradient for my parameters).</p>
</li>
<li>
<p>Grab the correct loss gradients from my cache and overwrite the incorrect gradients (from my auxiliary loss function) with the correct gradients from my loss function, so I effectively solve the problem of needing the backward hooks to return a different <code>grad_output</code> value (than the one from my loss function).</p>
</li>
<li>
<p>(To summarize in-case that was confusing!) I want the gradient from my loss function but want the <code>grad_output</code> values (which comes from the <code>full_backward_hook</code>) from a different loss function.</p>
</li>
</ul>
<p>I’ve managed to do this with the caching procedure I’ve described above. However, it only works for one epoch. Which is odd, because PyTorch is telling me from the Traceback that it failed to set up the hooks but it’s clearly done it from the previous epoch? The Full traceback is as follows.</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "test_hooks_double_loss.py", line 120, in &lt;module&gt;
    X, acceptance = sampler(burn_in)
  File "~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1065, in _call_impl
    return forward_call(*input, **kwargs)
  File "~/.local/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 28, in decorate_context
    return func(*args, **kwargs)
  File "~/Code/Samplers.py", line 94, in forward
    self.step()
  File "~/.local/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 28, in decorate_context
    return func(*args, **kwargs)
  File "~/Code/Samplers.py", line 72, in step
    log_pdf_chains = self._log_pdf(self.chains).detach_()
  File "~/.local/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 28, in decorate_context
    return func(*args, **kwargs)
  File "~/Code/Samplers.py", line 48, in _log_pdf
    return self.network(x)[1].mul(2).detach_()
  File "~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1065, in _call_impl
    return forward_call(*input, **kwargs)
  File "~/Code/Models.py", line 59, in forward
    log_envs = self.log_envelope(x0)
  File "~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1081, in _call_impl
    input = bw_hook.setup_input_hook(input)
  File "~/.local/lib/python3.8/site-packages/torch/utils/hooks.py", line 158, in setup_input_hook
    res, input_idx = self._apply_on_tensors(fn, args)
  File "~/.local/lib/python3.8/site-packages/torch/utils/hooks.py", line 142, in _apply_on_tensors
    raise RuntimeError("Error while setting up backward hooks. Please open "
RuntimeError: Error while setting up backward hooks. Please open an issue with a code sample to reproduce this.

</code></pre>
<p>Although, the Traceback states to make a reproducible example it seems quite difficultt given I’ve probably got near 1000 lines of code over numerous files. Would sharing perhaps parts of <code>main.py</code> with my caching be an acceptable start to debugging this?</p>
<p>Any help with this would be greatly appreciated! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>Thank you!</p>
<p>Edit: I assume there’s something wrong with <code>self.log_envelope(x0)</code>? The source code for that layer is given below,</p>
<pre><code class="lang-auto">class LogEnvelope(nn.Module):

  def __init__(self, num_particles, num_dets, bias=False) -&gt; None:
    super(LogEnvelope, self).__init__()
    
    self.num_particles = num_particles
    self.num_dets = num_dets
    
    self.weight = nn.Parameter(torch.empty(self.num_dets, self.num_particles))
    self.reset_parameters()

  def reset_parameters(self) -&gt; None:
    torch.nn.init.uniform_(self.weight, 1., 2.) 

  def forward(self, x0):
    return torch.einsum("bin,di-&gt;bdin", x0, self.weight).pow(2).sum(dim=-1)  
</code></pre>

ptrblck: <p>Thanks for reporting this! Could you create an issue on <a href="https://github.com/pytorch/pytorch/issues">GitHub</a> with your code snippet (even if it’s long) to get the visibility of the code owners, please?</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>! I can do, I’ll try and reduce it into a slimmer file (and hopefully just 1 file, although I can’t promise anything!)  and open a new issue on github! Thanks!</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>,</p>
<p>During the process of creating a standalone script, I found the error and it truly is bizarre. The problem emerges from the file <code>Samplers.py</code>. It’s effectively a module that uses a Network to calculate input data in accordance to a given probability by utilizing the Metropolis-Hastings algorithm and a random walk.</p>
<p>In short, and to add some context, it generates input data for my Feed-Forward Neural Network (Which is an R^N → R^1 function) by creating input data which is distributed to the output of the network squared.</p>
<pre><code class="lang-auto">input_data, acceptance  = Sampler(10) #10 steps of Metropolis-Hastings (input_data is shape [B,N])
                                      #input_data is distributed according to net(input_data).pow(2)

output = net(input_data) #my R^N -&gt; R^1 function 
</code></pre>
<p>The line of code that caused the crash is shown below,</p>
<pre><code class="lang-auto">  @torch.no_grad()
  def forward(self, burn_in: int) -&gt; Tuple[Tensor, Tensor]:
    self.acceptance.zero_()
    
    for i in range(burn_in+1): #check +1 here
      self.step()
    #return self.chains.detach_().requires_grad_(), self.acceptance.detach_() #crashes
    return self.chains.detach_(), self.acceptance.detach_() #doesn't crash
</code></pre>
<p>I called <code>requires_grad_()</code> on the data that’s returned from the <code>Sampler</code> as I need a gradient for it within my loss function. I do enable <code>requires_grad_()</code> within my loss function to check it’s enabled, so I can remove <code>requires_grad_()</code> within the <code>Sampler</code> and still calculate my loss. But to me it seems weird that this usage causes such a weird error.</p>
<p>Thank you!</p>

ptrblck: <p>Thanks for the update!<br>
Do you require gradients in the samples themselves or are you also training the network used in the sampler?<br>
Also, your current <code>forward</code> method is using the <code>no_grad</code> context manager. I assume the <code>forward</code> is used in the sampler network?<br>
If so, does the code also raise the error, if you call <code>detach_().requires_grad_()</code> outside of the <code>forward</code> on the returned samples?</p>

AlphaBetaGamma96: <aside class="quote no-group" data-username="ptrblck" data-post="5" data-topic="133302">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>Do you require gradients in the samples themselves or are you also training the network used in the sampler?</p>
</blockquote>
</aside>
<p>I’m training the network within the sampler. (I do require gradients for the samples, but it’s a 2nd order derivative and I’m following this neat trick here which reshapes the input and ensure they require_grads)</p><aside class="quote" data-post="3" data-topic="41667">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/j/a8b319/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/how-to-calculate-laplacian-sum-of-2nd-derivatives-in-one-step/41667/3">How to calculate Laplacian (sum of 2nd derivatives) in one step</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>
    Can you use the new Hessian function introduced in 1.5? <a href="https://pytorch.org/docs/stable/autograd.html#torch.autograd.functional.hessian" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/autograd.html#torch.autograd.functional.hessian</a>
  </blockquote>
</aside>

<aside class="quote no-group" data-username="ptrblck" data-post="5" data-topic="133302">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>Also, your current <code>forward</code> method is using the <code>no_grad</code> context manager. I assume the <code>forward</code> is used in the sampler network?</p>
</blockquote>
</aside>
<p>Yes, it is used. This method creates my input data, which to give it a more mathematical context rather than purely computational, are a set of positions within some N-dimensional vector space. These positions are initially randomly distributed (via <code>torch.randn</code> then via the <a href="https://en.wikipedia.org/wiki/Metropolis%E2%80%93Hastings_algorithm" rel="noopener nofollow ugc">Metropolis-Hastings algorthim</a> are distrubted according to my probability distribution. I don’t require any gradient computations here, only function evaluations so I decided to run the forward method with a <code>no_grad</code> context manager in order to 1) get a speed-up and 2) not backprop through the sampling algorithm when I call my loss function!</p>
<p>Does that make things clearer as to why I implemented the <code>forward</code> method like it that way? (If there’s a mistake, please let me know!)</p>
<aside class="quote no-group" data-username="ptrblck" data-post="5" data-topic="133302">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>If so, does the code also raise the error, if you call <code>detach_().requires_grad_()</code> outside of the <code>forward</code> on the returned samples?</p>
</blockquote>
</aside>
<p>Good idea. I’ll have a look and follow up with this!</p>

AlphaBetaGamma96: <aside class="quote no-group" data-username="ptrblck" data-post="5" data-topic="133302">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>If so, does the code also raise the error, if you call <code>detach_().requires_grad_()</code> outside of the <code>forward</code> on the returned samples?</p>
</blockquote>
</aside>
<p>So I moved the <code>detach_().requires_grad_()</code> outside the <code>forward</code> method and it crashes like before. I’ve tried a few different ways of doing this and I think I have an idea as to what’s causing it.</p>
<pre><code class="lang-auto">@torch.no_grad()
def forward(self, burn_in: int) -&gt; Tuple[Tensor, Tensor]:
  self.acceptance.zero_()
    
  for i in range(burn_in+1): #check +1 here
    self.step()
  #return self.chains.detach_().requires_grad_(), self.acceptance.detach_() #doesn't work
  return self.chains.detach_(), self.acceptance.detach_() #works
  return self.chains.detach().requires_grad_(), self.acceptance.detach_() #works
</code></pre>
<p>So it seems perhaps the issue arises from calling an in-place op on an in-place op?</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a> ,</p>
<p>I managed to make a reproducible example in less than 25 lines <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=10" title=":smiley:" class="emoji" alt=":smiley:"></p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

#The same as my metropolis-hasting algorthim (albeit shorter)
x = torch.randn(4,2,requires_grad=True)
fc = nn.Linear(2, 1)
x.detach_().requires_grad_()

def _save_output(module, grad_input, grad_output) -&gt; None:
  print("grad_output[0]: ",grad_output[0])

fc.register_full_backward_hook(_save_output)

with torch.no_grad():
  y = fc(x)

z = x.pow(2).sum(dim=-1)

#loss is defined as the sum of 2 terms
#first term, y has no gradient to it (purely scales)
#second term, z has a gradient. (in short, y scales z)
loss = torch.sum(y*z)

loss.backward()
</code></pre>
<p>when running this example, it returns,</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "backward_hook_setup_error_standalone.py", line 15, in &lt;module&gt;
    y = fc(x)
  File "~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1081, in _call_impl
    input = bw_hook.setup_input_hook(input)
  File "~/.local/lib/python3.8/site-packages/torch/utils/hooks.py", line 158, in setup_input_hook
    res, input_idx = self._apply_on_tensors(fn, args)
  File "~/.local/lib/python3.8/site-packages/torch/utils/hooks.py", line 142, in _apply_on_tensors
    raise RuntimeError("Error while setting up backward hooks. Please open "
RuntimeError: Error while setting up backward hooks. Please open an issue with a code sample to reproduce this.
</code></pre>
<p>Shall I open an issue with this?</p>

ptrblck: <p>Great! Yes, please create an issue using this code snippet! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

stas: <p>For posterity - this bug has been fixed in <code>torch==1.10</code></p>
<aside class="onebox githubissue" data-onebox-src="https://github.com/pytorch/pytorch/issues/66067">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/issues/66067" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewbox="0 0 14 16" aria-hidden="true"><path d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/66067" target="_blank" rel="noopener nofollow ugc">RuntimeError: Error while setting up backward hooks. Please open an issue with a code sample to reproduce this </a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2021-10-04" data-time="08:51:05" data-timezone="UTC">08:51AM - 04 Oct 21 UTC</span>
      </div>

        <div class="date">
          closed <span class="discourse-local-date" data-format="ll" data-date="2021-10-05" data-time="21:38:44" data-timezone="UTC">09:38PM - 05 Oct 21 UTC</span>
        </div>

      <div class="user">
        <a href="https://github.com/AlphaBetaGamma96" target="_blank" rel="noopener nofollow ugc">
          <img alt="AlphaBetaGamma96" src="https://avatars.githubusercontent.com/u/69877955?v=4" class="onebox-avatar-inline" width="20" height="20">
          AlphaBetaGamma96
        </a>
      </div>
    </div>

    <div class="labels">
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          module: nn
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          triaged
        </span>
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">## 🐛 Bug



Setting up backward hooks fails when the loss function depends o<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">n a detached Tensor. 

## To Reproduce
```
import torch
import torch.nn as nn

x = torch.randn(4,2,requires_grad=True)
fc = nn.Linear(2, 1)

#x.detach_().requires_grad_()

def _save_output(module, grad_input, grad_output) -&gt; None:
  print("grad_output[0]: ",grad_output[0])

fc.register_full_backward_hook(_save_output)

with torch.no_grad():
  y = fc(x)

z = x.pow(2).sum(dim=-1)

#loss is defined as the sum of 2 terms
#first term, y has no gradient to it (purely scales)
#second term, z has a gradient. (in short, y scales z)
loss = torch.sum(y+z)

loss.backward()
```


```
Traceback (most recent call last):
  File "backward_hook_setup_error_standalone.py", line 15, in &lt;module&gt;
    y = fc(x)
  File "~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1081, in _call_impl
    input = bw_hook.setup_input_hook(input)
  File "~/.local/lib/python3.8/site-packages/torch/utils/hooks.py", line 158, in setup_input_hook
    res, input_idx = self._apply_on_tensors(fn, args)
  File "~/.local/lib/python3.8/site-packages/torch/utils/hooks.py", line 142, in _apply_on_tensors
    raise RuntimeError("Error while setting up backward hooks. Please open "
RuntimeError: Error while setting up backward hooks. Please open an issue with a code sample to reproduce this.
```
## Expected behavior

The ability to set up backward hooks on layers where the loss partially contains detached Tensors.

## Environment
```
Collecting environment information...
PyTorch version: 1.10.0.dev20210824+cu111
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.3 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31

Python version: 3.8.8 (default, Apr 13 2021, 19:58:26)  [GCC 7.3.0] (64-bit runtime)
Python platform: Linux-5.11.0-37-generic-x86_64-with-glibc2.10
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: GeForce RTX 2060
Nvidia driver version: 460.91.03
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] functorch==0.0.1a0+0ae60dd
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.21.2
[pip3] numpydoc==1.1.0
[pip3] pytorch-lightning==1.3.5
[pip3] torch==1.10.0.dev20210824+cu111
[pip3] torch-kfac==0.0.1
[pip3] torchaudio==0.9.0
[pip3] torchdiffeq==0.2.2
[pip3] torchdyn==0.2.2.1
[pip3] torchmetrics==0.3.2
[pip3] torchsde==0.2.5
[pip3] torchvision==0.10.0+cu111
[pip3] torchviz==0.0.2
[conda] blas                      1.0                         mkl  
[conda] mkl                       2021.2.0           h06a4308_296  
[conda] mkl-service               2.3.0            py38h27cfd23_1  
[conda] mkl_fft                   1.3.0            py38h42c9631_2  
[conda] mkl_random                1.2.1            py38ha9443f7_2  
[conda] mypy_extensions           0.4.3                    py38_0  
[conda] numpy                     1.19.5                   pypi_0    pypi
[conda] numpy-base                1.20.1           py38h7d8b39e_0  
[conda] numpydoc                  1.1.0              pyhd3eb1b0_1  
[conda] pytorch-lightning         1.3.5                    pypi_0    pypi
[conda] torch                     1.9.0+cu111              pypi_0    pypi
[conda] torchaudio                0.9.0                    pypi_0    pypi
[conda] torchdiffeq               0.2.2                    pypi_0    pypi
[conda] torchdyn                  0.2.2.1                  pypi_0    pypi
[conda] torchmetrics              0.3.2                    pypi_0    pypi
[conda] torchsde                  0.2.5                    pypi_0    pypi
[conda] torchvision               0.10.0+cu111             pypi_0    pypi
[conda] torchviz                  0.0.2                    pypi_0    pypi
```

## Additional context


This was originally opened as a thread on the forums, the link to it is here: https://discuss.pytorch.org/t/runtimeerror-error-while-setting-up-backward-hooks-please-open-an-issue-with-a-code-sample-to-reproduce-this/133302

cc @albanD @mruberry @jbschlosser @walterddr</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>So you want:</p>
<pre><code class="lang-auto">import torch
from packaging import version
[...]
    def _register_backward_hook(self, module):
        if version.parse(torch.__version__) &gt;= version.parse("1.10"):
            module.register_full_backward_hook(self.backward_hook)
        else:
            module.register_backward_hook(self.backward_hook)
</code></pre>