Tracking down NaN gradients

joel: <p>I have noticed that there are NaNs in the gradients of my model. This is confirmed by <code>torch.autograd.detect_anomaly()</code>:</p>
<p><code>RuntimeError: Function 'DivBackward0' returned nan values in its 1th output.</code></p>
<p>I do not know which division causes the problem since DivBackward0 does not seem to be a unique name. However, I have added asserts to all divisions (like <code>assert torch.all(divisor != 0)</code>) and also have lots of asserts to check for NaNs in general (like <code>assert torch.all(~torch.isnan(t))</code>).</p>
<p>I also iterate the graph and register hooks which print the function and check for NaNs with the following code:</p>
<pre><code class="lang-auto">def iter_graph(root, callback):
    queue = [root]
    seen = set()
    while queue:
        fn = queue.pop()
        if fn in seen:
            continue
        seen.add(fn)
        for next_fn, _ in fn.next_functions:
            if next_fn is not None:
                queue.append(next_fn)
        callback(fn)

def register_hooks(var):
    fn_dict = {}
    def hook_cb(fn):
        def register_grad(grad_input, grad_output):
            print(fn)
            assert all(t is None or torch.all(~torch.isnan(t)) for t in grad_input), f"{fn} grad_input={grad_input} grad_output={grad_output}"
            assert all(t is None or torch.all(~torch.isnan(t)) for t in grad_output), f"{fn} grad_input={grad_input} grad_output={grad_output}"
            
            fn_dict[fn] = grad_input
        fn.register_hook(register_grad)
    iter_graph(var.grad_fn, hook_cb)
</code></pre>
<p>The output looks like this:</p>
<pre><code class="lang-auto">&lt;ViewBackward object at 0x7fb79bae50d0&gt;
&lt;SubBackward0 object at 0x7fb79bae5130&gt;
&lt;DivBackward0 object at 0x7fb79bae51c0&gt;
&lt;DivBackward0 object at 0x7fb79bae5190&gt;
&lt;SliceBackward object at 0x7fb79bae50a0&gt;
&lt;SliceBackward object at 0x7fb79bae5400&gt;
&lt;ViewBackward object at 0x7fb79badcfd0&gt;
...
&lt;SigmoidBackward object at 0x7fb79bacc3d0&gt;
&lt;AddmmBackward object at 0x7fb79bacc430&gt;
&lt;TBackward object at 0x7fb79bacc4c0&gt;
&lt;CudnnBatchNormBackward object at 0x7fb79bacc490&gt;
...
&lt;torch.autograd.function.BilinearInterpolationBackward object at 0x7fb8cc79c4a0&gt;
&lt;torch.autograd.function.BilinearInterpolationBackward object at 0x7fb8cc79c3c0&gt;
&lt;torch.autograd.function.BilinearInterpolationBackward object at 0x7fb8cc79c2e0&gt;
</code></pre>
<p>And then it fails with:</p>
<pre><code class="lang-auto">AssertionError: &lt;torch.autograd.function.BilinearInterpolationBackward object at 0x7fb8cc79c2e0&gt; grad_input=(tensor([[0.],
        [0.],
        [0.],
        [0.],
        [0.],
        [nan],
        [0.],
        [0.],
        [0.],
        [nan]], device='cuda:0'), tensor([[0.],
        [0.],
        [0.],
        [0.],
        [0.],
        [nan],
        [0.],
        [0.],
        [0.],
        [nan]], device='cuda:0'), None, None, None, None, None, None, None, None, None) grad_output=(tensor([[ 0.0000e+00, -7.8456e+29],
        [ 0.0000e+00,  2.4914e+31],
        [ 0.0000e+00,  0.0000e+00],
        [ 0.0000e+00, -2.4474e+30],
        [ 0.0000e+00,  5.9677e+30],
        [        nan,  0.0000e+00],
        [ 0.0000e+00,  0.0000e+00],
        [ 0.0000e+00,  9.7542e+30],
        [ 0.0000e+00, -2.9419e+30],
        [        nan,  0.0000e+00]], device='cuda:0'),)
</code></pre>
<p>Interestingly, it does not fail immediately after <code>DivBackward0</code>. <code>BilinearInterpolationBackward</code> has NaNs in both the inputs and the output, which means that it does not cause the problem either.</p>
<p>I am pretty lost at this point. What else can I do to track down the NaN gradients?</p>
<p><strong>Edit:</strong></p>
<ul>
<li>Checking for Inf does not help</li>
<li>Bigger batch size does not help</li>
</ul>
<p><strong>Edit 2:</strong></p>
<p>If I disable cuDNN with <code>torch.backends.cudnn.enabled = False</code>, then I get infinity in a <code>MulBackward0</code>. Investigating further.</p>

albanD: <p>Hi,</p>
<p>Didn’t the second stack trace when the anomaly mode is enabled show you where the bad div is?</p>

joel: <p>No, the stack trace is useless in this case. The exception is thrown in C++. Also, by now I am pretty sure that div isn’t the problem here. I guess detect_anomaly does not work well with some cudnn optimizations.</p>

albanD: <p>Can you still share the two stack traces?</p>

joel: <p>This is the one with cudnn disabled:</p>
<pre><code class="lang-auto">---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-5-3906501279d7&gt; in &lt;module&gt;
     21 
     22         # get_dot = register_hooks(loss)
---&gt; 23         loss.backward()
     24         # dot = get_dot()
     25 

venv/lib/python3.8/site-packages/torch/tensor.py in backward(self, gradient, retain_graph, create_graph)
    193                 products. Defaults to ``False``.
    194         """
--&gt; 195         torch.autograd.backward(self, gradient, retain_graph, create_graph)
    196 
    197     def register_hook(self, hook):

venv/lib/python3.8/site-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables)
     95         retain_graph = create_graph
     96 
---&gt; 97     Variable._execution_engine.run_backward(
     98         tensors, grad_tensors, retain_graph, create_graph,
     99         allow_unreachable=True)  # allow_unreachable flag

RuntimeError: Function 'SelectBackward' returned nan values in its 0th output.
</code></pre>
<p>And with cudnn:</p>
<pre><code class="lang-auto">---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-5-3906501279d7&gt; in &lt;module&gt;
     21 
     22         # get_dot = register_hooks(loss)
---&gt; 23         loss.backward()
     24         # dot = get_dot()
     25 

venv/lib/python3.8/site-packages/torch/tensor.py in backward(self, gradient, retain_graph, create_graph)
    193                 products. Defaults to ``False``.
    194         """
--&gt; 195         torch.autograd.backward(self, gradient, retain_graph, create_graph)
    196 
    197     def register_hook(self, hook):

venv/lib/python3.8/site-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables)
     95         retain_graph = create_graph
     96 
---&gt; 97     Variable._execution_engine.run_backward(
     98         tensors, grad_tensors, retain_graph, create_graph,
     99         allow_unreachable=True)  # allow_unreachable flag

RuntimeError: Function 'DivBackward0' returned nan values in its 1th output.
</code></pre>

albanD: <p>Isn’t there a second stack trace just above this one when anomaly mode is enabled?</p>

joel: <p>No. Should there be one? I am using PyTorch 1.4.</p>

albanD: <p>For example in the code below, you can see the warning that points to the faulty forward part:</p>
<pre><code class="lang-python">import torch
from torch import autograd
autograd.set_detect_anomaly(True)

return_nan = True

class MyBadFn(autograd.Function):
    @staticmethod
    def forward(ctx, inp):
        ctx.save_for_backward(inp)
        return inp ** 2

    @staticmethod
    def backward(ctx, grad_out):
        inp, = ctx.saved_tensors
        res = 2 * grad_out * inp
        if return_nan:
            res = grad_out / 0 * 0
        return res


inp = torch.rand(10, requires_grad=True)

out = MyBadFn.apply(inp)

out.sum().backward()
</code></pre>
<p>prints</p>
<pre><code class="lang-auto">[W python_anomaly_mode.cpp:60] Warning: Error detected in MyBadFnBackward. Traceback of forward call that caused the error:
  File "foo.py", line 24, in &lt;module&gt;
    out = MyBadFn.apply(inp)
 (function print_stack)
Traceback (most recent call last):
  File "foo.py", line 26, in &lt;module&gt;
    out.sum().backward()
  File "/Users/albandes/workspace/pytorch_dev/torch/tensor.py", line 184, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/Users/albandes/workspace/pytorch_dev/torch/autograd/__init__.py", line 107, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: Function 'MyBadFnBackward' returned nan values in its 0th output.
</code></pre>

joel: <p>For me this looks different:</p>
<pre><code class="lang-auto">---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-2-c48dcde9ba4c&gt; in &lt;module&gt;
     24 out = MyBadFn.apply(inp)
     25 
---&gt; 26 out.sum().backward()

venv/lib/python3.8/site-packages/torch/tensor.py in backward(self, gradient, retain_graph, create_graph)
    193                 products. Defaults to ``False``.
    194         """
--&gt; 195         torch.autograd.backward(self, gradient, retain_graph, create_graph)
    196 
    197     def register_hook(self, hook):

venv/lib/python3.8/site-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables)
     95         retain_graph = create_graph
     96 
---&gt; 97     Variable._execution_engine.run_backward(
     98         tensors, grad_tensors, retain_graph, create_graph,
     99         allow_unreachable=True)  # allow_unreachable flag

RuntimeError: Function 'MyBadFnBackward' returned nan values in its 0th output.
</code></pre>
<p>Regarding my “real” problem: I seem to have gradient explosion and am now trying to figure out why.</p>
<p><strong>Edit:</strong> I think my Jupyter Notebook is messing with the errors. When I run this in a normal Python file, I get the following instead:</p>
<pre><code class="lang-auto">venv/bin/python test.py
Warning: Traceback of forward call that caused the error:
  File "test.py", line 24, in &lt;module&gt;
    out = MyBadFn.apply(inp)
 (print_stack at /pytorch/torch/csrc/autograd/python_anomaly_mode.cpp:57)
Traceback (most recent call last):
  File "test.py", line 26, in &lt;module&gt;
    out.sum().backward()
  File "venv/lib/python3.8/site-packages/torch/tensor.py", line 195, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "venv/lib/python3.8/site-packages/torch/autograd/__init__.py", line 97, in backward
    Variable._execution_engine.run_backward(
RuntimeError: Function 'MyBadFnBackward' returned nan values in its 0th output.
</code></pre>

albanD: <p>Ho <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=9" title=":confused:" class="emoji" alt=":confused:"> Is the Jupyter Notebook silencing warnings?? (because the second stack trace is actually a python warning), what happens if you do: <code>import warning; warnings.simplefilter("always")</code> in the notebook to try and make all warnings appear?</p>

joel: <p>This does not help either. In case you want to reproduce:<br>
Python 3.8.2<br>
torch==1.4.0</p>
<pre><code class="lang-auto">$ jupyter --version
jupyter core     : 4.6.3
jupyter-notebook : 6.0.3
qtconsole        : 4.7.2
ipython          : 7.13.0
ipykernel        : 5.2.0
jupyter client   : 6.1.2
jupyter lab      : 2.0.1
nbconvert        : 5.6.1
ipywidgets       : 7.5.1
nbformat         : 5.0.5
traitlets        : 4.3.3
</code></pre>
<p>I have tired also some other stuff and did not get it working.</p>

albanD: <p>Interesting, I can actually repro on colab: <a href="https://colab.research.google.com/drive/1gbUpNxLMlL4H-mKsgTh7UiU_JTSPHnvZ">https://colab.research.google.com/drive/1gbUpNxLMlL4H-mKsgTh7UiU_JTSPHnvZ</a><br>
I’ll open an issue to check that.</p>
<p>Going back to your original issue, if you can run the same code in a scirpt outside of the notebook, that might help you pinpoint things.<br>
Otherwise, if you can make a simple repro, I can take a closer look as well.</p>

joel: <p>Thanks for your help and creating the issue. I have been able to solve the problem since. I really had a gradient explosion, which caused Infs (which were turned into NaNs by cudnn).</p>
<p>To mitigate the problem, I am now clipping the gradients in a <code>backward</code> of a <code>torch.autograd.Function</code> (<code>register_hook</code> would work as well). I cannot use normal gradient clipping, since it happens immediately in the first iteration.</p>

aswamy: <p>Hello <a class="mention" href="/u/joel">@joel</a> <a class="mention" href="/u/alband">@albanD</a>,</p>
<p>I am also encountering the same problem in my code. In my case, MSE loss is nan. I tried clipping gradients by adding below lines</p>
<p>torch.nn.utils.clip_grad_norm_(self.net.parameters(), max_norm=1.0, norm_type=2.0)<br>
optimizer.step()<br>
optimizer.zero_grad()</p>
<p>But I still get nan, would you be able to share what actually you modified?</p>

Andy_Zhu: <p>i get the same error too, I simplify my net to one Linear, and remove the activation function, and this error disappeared, but I got another error:Function ‘AddmmBackward’ returned nan values in its 1th output.</p>

Vignesh_Kannan: <p>The notebook doesn’t seem to exist anymore, <a class="mention" href="/u/alband">@albanD</a> - do you have another version for the same?</p>

albanD: <p>The notebook was just running the code in the message above and showing that the warning was silenced.</p>

Adex: <aside class="quote no-group" data-username="joel" data-post="13" data-topic="78112">
<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/joel/48/13872_2.png" class="avatar"> joel:</div>
<blockquote>
<p>torch.autograd.Function</p>
</blockquote>
</aside>
<p>Can you please share details about your approach? I am also facing the same problem and normal gradient clipping is of no use.</p>