Grad becomes nan for all parameters

JiarunLiu: <p>I’m using <a href="https://github.com/facebookresearch/mae" rel="noopener nofollow ugc">MAE</a> to pretrain a ViT model on my custom dataset with 4 A800 GPU.</p>
<p>I found that all gradients are nan after epoch 486. (The grad here is manually saved and printed)<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/0/f03f33260ed7b42287dc4ea125aea2ddec63cd06.png" data-download-href="https://discuss.pytorch.org/uploads/default/f03f33260ed7b42287dc4ea125aea2ddec63cd06" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/0/f03f33260ed7b42287dc4ea125aea2ddec63cd06_2_690x261.png" alt="image" data-base62-sha1="yhjZRKsRpSKnQpXNQtvauTdZ61E" width="690" height="261" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/0/f03f33260ed7b42287dc4ea125aea2ddec63cd06_2_690x261.png, https://discuss.pytorch.org/uploads/default/original/3X/f/0/f03f33260ed7b42287dc4ea125aea2ddec63cd06.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/f/0/f03f33260ed7b42287dc4ea125aea2ddec63cd06.png 2x" data-dominant-color="323333"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">938×356 48.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>There loss looks good during the triaining, no nan or inf in the loss. But the model’s parameters won’t update anymore.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/0/204cc4fb85be59514c2ee36021123ff6f3392f57.png" data-download-href="https://discuss.pytorch.org/uploads/default/204cc4fb85be59514c2ee36021123ff6f3392f57" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/0/204cc4fb85be59514c2ee36021123ff6f3392f57_2_690x369.png" alt="image" data-base62-sha1="4BJKVplkGf1R0wpxgqaZwjeUTSD" width="690" height="369" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/0/204cc4fb85be59514c2ee36021123ff6f3392f57_2_690x369.png, https://discuss.pytorch.org/uploads/default/optimized/3X/2/0/204cc4fb85be59514c2ee36021123ff6f3392f57_2_1035x553.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/2/0/204cc4fb85be59514c2ee36021123ff6f3392f57_2_1380x738.png 2x" data-dominant-color="FCF7EA"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1850×990 44 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I tried to use  <code>torch.autograd.detect_anomaly()</code> to figure out where the issue comes from:</p>
<pre><code class="lang-auto">/usr/local/anaconda3/envs/jiarun_mae/lib/python3.8/site-packages/torch/autograd/__init__.py:197: UserWarning: Error detected in MulBackward0. Traceback of forward call that caused the error:
  File "main_pretrain.py", line 229, in &lt;module&gt;
    main(args)
  File "main_pretrain.py", line 198, in main
    train_stats = train_one_epoch(
  File "/home/jiarunliu/Documents/scaling/mae/engine_pretrain.py", line 69, in train_one_epoch
    loss, _, _ = model(samples, mask_ratio=args.mask_ratio)
  File "/usr/local/anaconda3/envs/jiarun_mae/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/anaconda3/envs/jiarun_mae/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 1040, in forward
    output = self._run_ddp_forward(*inputs, **kwargs)
  File "/usr/local/anaconda3/envs/jiarun_mae/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 1000, in _run_ddp_forward
    return module_to_run(*inputs[0], **kwargs[0])
  File "/usr/local/anaconda3/envs/jiarun_mae/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/jiarunliu/Documents/scaling/mae/models_mae.py", line 229, in forward
    loss = self.forward_loss(imgs, pred, mask)
  File "/home/jiarunliu/Documents/scaling/mae/models_mae.py", line 221, in forward_loss
    loss = (loss * mask).sum() / mask.sum()  # mean loss on removed patches
  File "/usr/local/anaconda3/envs/jiarun_mae/lib/python3.8/site-packages/torch/fx/traceback.py", line 57, in format_stack
    return traceback.format_stack()
 (Triggered internally at ../torch/csrc/autograd/python_anomaly_mode.cpp:114.)


Traceback (most recent call last):
  File "main_pretrain.py", line 229, in &lt;module&gt;
    main(args)
  File "main_pretrain.py", line 198, in main
    train_stats = train_one_epoch(
  File "/home/jiarunliu/Documents/scaling/mae/engine_pretrain.py", line 90, in train_one_epoch
    loss_scaler(loss, optimizer, 
  File "/home/jiarunliu/Documents/scaling/mae/util/misc.py", line 302, in __call__
    self._scaler.scale(loss).backward(create_graph=create_graph)
  File "/usr/local/anaconda3/envs/jiarun_mae/lib/python3.8/site-packages/torch/_tensor.py", line 487, in backward
    torch.autograd.backward(
  File "/usr/local/anaconda3/envs/jiarun_mae/lib/python3.8/site-packages/torch/autograd/__init__.py", line 197, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: Function 'MulBackward0' returned nan values in its 0th output.
</code></pre>
<p>The corresponding code in the traceback is:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/facebookresearch/mae/blob/efb2a8062c206524e35e47d04501ed4f544c0ae8/models_mae.py#L198-L214">
  <header class="source">

      <a href="https://github.com/facebookresearch/mae/blob/efb2a8062c206524e35e47d04501ed4f544c0ae8/models_mae.py#L198-L214" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/facebookresearch/mae/blob/efb2a8062c206524e35e47d04501ed4f544c0ae8/models_mae.py#L198-L214" target="_blank" rel="noopener nofollow ugc">facebookresearch/mae/blob/efb2a8062c206524e35e47d04501ed4f544c0ae8/models_mae.py#L198-L214</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="198" style="counter-reset: li-counter 197 ;">
          <li>def forward_loss(self, imgs, pred, mask):</li>
          <li>    """</li>
          <li>    imgs: [N, 3, H, W]</li>
          <li>    pred: [N, L, p*p*3]</li>
          <li>    mask: [N, L], 0 is keep, 1 is remove, </li>
          <li>    """</li>
          <li>    target = self.patchify(imgs)</li>
          <li>    if self.norm_pix_loss:</li>
          <li>        mean = target.mean(dim=-1, keepdim=True)</li>
          <li>        var = target.var(dim=-1, keepdim=True)</li>
          <li>        target = (target - mean) / (var + 1.e-6)**.5</li>
          <li></li>
          <li>    loss = (pred - target) ** 2</li>
          <li>    loss = loss.mean(dim=-1)  # [N, L], mean loss per patch</li>
          <li></li>
          <li>    loss = (loss * mask).sum() / mask.sum()  # mean loss on removed patches</li>
          <li>    return loss</li>
      </ol>
    </code></pre>



  </article>

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

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

<p>Besides, I use <code>amp + bf16</code> during training, otherwise the loss will become nan at about epoch 380.</p>
<p>I have tried following strategies but none of them works:</p>
<ul>
<li>reduce/increase learning rate</li>
<li>clip gradient by 1.0 or 0.1</li>
<li>replace the AdamW optimizer with Adam</li>
<li>enable/disable norm_pix_loss</li>
<li>use amp+bf16</li>
<li>use fp32</li>
</ul>
<p>Additionally, use a larger learning rate will produce <code>inf</code> in gradients:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/a/1afd15fba4b18eba94e5b9bce7935f267cf4c6cd.png" data-download-href="https://discuss.pytorch.org/uploads/default/1afd15fba4b18eba94e5b9bce7935f267cf4c6cd" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/a/1afd15fba4b18eba94e5b9bce7935f267cf4c6cd_2_690x396.png" alt="image" data-base62-sha1="3QKEbIuCM6cgMSiQH0vojfGvcp7" width="690" height="396" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/a/1afd15fba4b18eba94e5b9bce7935f267cf4c6cd_2_690x396.png, https://discuss.pytorch.org/uploads/default/original/3X/1/a/1afd15fba4b18eba94e5b9bce7935f267cf4c6cd.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/1/a/1afd15fba4b18eba94e5b9bce7935f267cf4c6cd.png 2x" data-dominant-color="3B3939"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">824×474 109 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Is there any possible reason for this issue? My package version is <code>torch==1.13.0+cu117</code>, <code>timm==0.3.2</code>.</p>
<p>Thank you!</p>

andrea_senz: <p>Hey <img src="https://discuss.pytorch.org/images/emoji/apple/wave.png?v=12" title=":wave:" class="emoji" alt=":wave:" loading="lazy" width="20" height="20">,<br>
Seems to be an exploding gradient issue.<br>
Did you standardize/normalize your input over rgb channels?<br>
Hope it helps.</p>

JiarunLiu: <p>Thank you for your quick response!</p>
<p>Yes, I have normalized my input data and set <code>self.norm_pix_loss = True</code>. Here is my data augmentation strategy, the mean and std are computed on my dataset.</p>
<pre><code class="lang-auto">    transform_train = transforms.Compose([
            transforms.RandomResizedCrop(args.input_size, scale=(0.2, 1.0), interpolation=3),  # 3 is bicubic
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.4271, 0.4054, 0.4118], std=[0.2124, 0.2165, 0.2112])])
</code></pre>

ptrblck: <p>Based on the stacktrace the <code>loss</code> calculation might cause the issues e.g. if <code>mask</code> contains only zeros:</p>
<pre data-code-wrap="python"><code class="lang-python">weight = torch.randn(10, 10, requires_grad=True)
loss = weight * 2
mask = torch.zeros_like(loss)
loss = (loss * mask).sum() / mask.sum()
print(loss)
# tensor(nan, grad_fn=&lt;DivBackward0&gt;)
loss.backward()
print(weight.grad)
# tensor([[nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
#         [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
#         [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
#         [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
#         [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
#         [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
#         [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
#         [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
#         [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
#         [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan]])
</code></pre>
<p>Did you make sure <code>mask</code> contains any valid values?</p>

JiarunLiu: <p>I have manually checked <code>mask.sum() == 150528</code> for each batch, so the loss should not be zero divided here. The weird thing is, my loss is a valid value all the time, although the grad is nan. The training can keep going but the network won’t actually update.</p>
<p>Besides, the most common traceback is:</p>
<pre><code class="lang-auto">RuntimeError: Function 'MulBackward0' returned nan values in its 0th output.
</code></pre>
<p>I’m not sure what this means? But it looks different from <code>DivBackward0</code>.</p>

JiarunLiu: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> I still not found the solution to solve this issue, but I have one successful attempt to train the model. I train the model with <code>amp(fp16)</code> precision for the first 300 epochs, then use <code>amp(bf16)</code> to train the following 500 epochs. Training with fp16 will have nan loss because some values will be out of the range of fp16, so I have to change the precision to bf16 or fp32 in the following epochs. Would my problem relate to the precision or pytorch version?</p>

ptrblck: <p>I doubt your issue is related to the PyTorch version, as the stacktrace did point to the lines of code I posted reproducing the issue and I also see <code>float16</code> being used in the MAE repository.</p>

JiarunLiu: <p>You are right, this issue no longer occurs when I upgrade torch from <code>1.13.0</code> to <code>2.2.2</code>. Thank you.</p>