RuntimeError: Function 'SoftmaxBackward' returned nan values in its 0th output

Alethia: <p>I’m training tacotron2 (a TTS model) using the seq2seq model with attention.<br>
I use nvidia apex to train the model with mixed precision, and I got the following error:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "train.py", line 709, in &lt;module&gt;
    main()
  File "train.py", line 650, in main
    scaled_loss.backward()
  File "/opt/conda/lib/python3.6/site-packages/torch/tensor.py", line 184, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/opt/conda/lib/python3.6/site-packages/torch/autograd/__init__.py", line 123, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: (140597717193280,)
</code></pre>
<p>This was confusing, so I added  <code>torch.autograd.set_detect_anomaly(True)</code> in the code. Then I got:</p>
<pre><code class="lang-auto">[W python_anomaly_mode.cpp:60] Warning: Error detected in SoftmaxBackward. Traceback of forward call that caused the error:
  File "train.py", line 709, in &lt;module&gt;
    main()
  File "train.py", line 618, in main
    args)
  File "/workspace/tacotron2/tacotron2/model.py", line 435, in decode
    attention_weights_cat, mask)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 577, in __call__
    result = self.forward(*input, **kwargs)
  File "/workspace/tacotron2/tacotron2/model.py", line 113, in forward
    attention_weights = F.softmax(alignment, dim=1)
  File "/opt/conda/lib/python3.6/site-packages/apex/amp/wrap.py", line 28, in wrapper
    return orig_fn(*new_args, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/functional.py", line 1498, in softmax
    ret = input.softmax(dim)
 (function print_stack)
Traceback (most recent call last):
  File "train.py", line 709, in &lt;module&gt;
    main()
  File "train.py", line 618, in main
    args)
  File "train.py", line 429, in consistency_training_subroutine
    scaled_loss.backward()
  File "/opt/conda/lib/python3.6/site-packages/torch/tensor.py", line 184, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/opt/conda/lib/python3.6/site-packages/torch/autograd/__init__.py", line 123, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: Function 'SoftmaxBackward' returned nan values in its 0th output.
</code></pre>
<p>Looks like the above error happened when computing the gradient of the <code>softmax</code> layer for computing attention weights.<br>
I first tried to print out the input tensor (including max and min value) of softmax, and I got:</p>
<pre><code class="lang-auto">tensor([[ 3.7246,  0.2996,  1.9951,  ...,  4.5039,  5.5117,  6.2266],
        [ 2.9258, -0.5249,  0.3879,  ...,  4.8828,  4.8750,  9.7031],
        [-6.9258, -8.2266, -8.1016,  ...,    -inf,    -inf,    -inf],
        ...,
        [-9.0312, -9.9531, -9.3047,  ...,    -inf,    -inf,    -inf],
        [-1.9766, -8.1562, -5.5586,  ...,    -inf,    -inf,    -inf],
        [-8.2656, -5.7539, -6.9531,  ...,    -inf,    -inf,    -inf]],
       device='cuda:0', dtype=torch.float16, grad_fn=&lt;MaskedFillBackward0&gt;)
tensor(12.4922, device='cuda:0', dtype=torch.float16, grad_fn=&lt;MaxBackward1&gt;)
tensor(-inf, device='cuda:0', dtype=torch.float16, grad_fn=&lt;MinBackward1&gt;)
</code></pre>
<p>Here the minimum <code>-inf</code> is filled for the padded positions in the batch.</p>
<p>I think this may because the maximal value <code>float16</code> can support is <code>65536</code>, <code>torch.exp(12.4922)</code> will output <code>inf</code>. So I closed the apex but the problem still appeared.</p>
<p>I wonder if there are good methods to solve this problem or help me understand this problem well.<br>
Thank you.</p>

ptrblck: <p>We recommend to use the native mixed-precision training utility via <code>torch.cuda.amp</code>, which will use <code>flaot32</code> for the <code>softmax</code> operation as given <a href="https://pytorch.org/docs/stable/amp.html#ops-that-can-autocast-to-float32">here</a>.</p>

divinho: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> I still get NaNs although surrounding the linear layer and log_softmax with <code>with torch.cuda.amp.autocast(enabled=False):</code></p>

ptrblck: <p>Check what the input <code>dtypes</code> to these operations are and make sure they are safe to use or just enable <code>auocast</code> as it will do the right casts for you.</p>