How to find causes of NAN grads?

username: mikasa
cooked: <p>After upgrade to PyTorch 1.12.1 (conda, cuda 11.3), I got <code>nan</code> grad for some <code>Conv2d</code> weights and biases right after the validation:</p>
<pre><code class="lang-auto">- Epoch 0 training.
- Validation.
- Epoch 1 training.
- Validation.
...
- Epoch N training.
- Validation.
- Got nan in the first step of epoch N+1.
- Got nan in the second step of epoch N+2.
- Got nan in the third step of epoch N+3.
...
</code></pre>
<p>If I save the model’s <code>state dict</code> and inputs when <code>nan</code> appears, then run forward and backward manually in another python process with the saved objects, the grad will be normal with no <code>nan</code>.</p>
<p>I use <code>model.to(memory_format=torch.channels_last)</code> and <code>fp16</code> training. If I remove <code>model.to(memory_format=torch.channels_last)</code> or use <code>fp32</code> training, the training will work normally without <code>nan</code> grad.</p>
<p>It also works in PyTorch 1.11.</p>
<p>Is there any suggestion to find causes of random <code>nan</code> grads?</p>

username: AlphaBetaGamma96
cooked: <p>Hi <a class="mention" href="/u/mikasa">@mikasa</a>,</p>
<p>Have a look at <code>torch.autograd.detect_analomy</code> which is designed to find what operations cause NaNs. (docs <a href="https://pytorch.org/docs/stable/autograd.html#torch.autograd.detect_anomaly" rel="noopener nofollow ugc">here</a>)</p>

username: mikasa
cooked: <p>Thank you <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a>. I have tried <code>torch.autograd.detect_anomaly</code>, but it just said <code>ConvolutionBackward0</code> returns <code>nan</code> values.</p>
<p>The convolution will not get <code>nan</code> values if I run the model in another python process with the same states and inputs. It also works normally if I use PyTorch 1.11 or use Pytorch 1.12 with data format <em>NCHW</em>.</p>

username: AlphaBetaGamma96
cooked: <aside class="quote no-group" data-username="mikasa" data-post="3" data-topic="160563">
<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/m/0ea827/48.png" class="avatar"> mikasa:</div>
<blockquote>
<p>ConvolutionBackward0</p>
</blockquote>
</aside>
<p>This isn’t referencing the output of your convolution layer, it’s referencing the gradients of the convolution layer for its 0-th output.</p>
<p>Could you share the model you’re using? As well as the full error message you get when running <code>torch.autograd.detect_anomlay</code> because if you have more than 1 convolution layer it’ll be harder to know which one is causing the NaN issue (although it’s most likely the last one).</p>

username: mikasa
cooked: <p>It seems that I can’t post the full model and trace here.<br>
I set the breakpoint at the line containing the conv with nan gradient and run the following code:</p>
<pre><code class="lang-python"># Conv2d.
print(module)  # Conv2d(32, 48, kernel_size=(8, 1), stride=(4, 1), padding=(2, 0), bias=False) 
print(module.weight.stride()) # (256, 1, 32, 32)
print(module.weight.dtype) # torch.float32
print(module.weight.isnan().any()) # torch(False, device='cuda:0')

# Conv2d's input.
print(input.shape) # torch.Size([48, 32, 64, 376])
print(input.stride()) # (770048, 1, 12032, 32)
print(input.dtype)  # torch.float16
print(input.isnan().any())  # torch(False, device='cuda:0')

# Ok.
print(torch.autograd.grad(module(input)[0, 0, 0, 0], input)[0].isnan().any()) # torch(False, device='cuda:0')
# Why does `* 0.0` get nan grad?
print(torch.autograd.grad(module(input)[0, 0, 0, 0] * 0.0, input)[0].isnan().any()) # torch(True, device='cuda:0')
</code></pre>
<p>I’m not sure why the last line (and the full conv backward) gets <code>nan</code> grad (in PyTorch 1.12.1) after running some validation passes.</p>