Apex: nan loss with O1 and O2

username: dhecloud
cooked: <p>Hi, i am using apex to try to fit a larger model (T5-large) on a single K40 GPU. i understand K40 is using the kepler gpu, so most of the benefits of apex will not be usable on the K40. However, i read that is it still possible to use fp16 to fit a larger model so i went ahead with that. I managed to start training with a larger model (T5-large) which otherwise would not be possible with the GPU’s specs.</p>
<p>My problem is that when when using opt-levels <code>O0</code> (fp32 training) the batch losses are fine (not nan). When i change to mixed precision training <code>O1</code>, <code>O2</code>, the batch losses start to output as <code>nan</code>. I also tried with smaller models (T5-base and T5-small)  where the same thing happened; training proceeded totally normally with <code>O0</code> (normal loss values) but had nan losses for <code>O1</code> and <code>O2</code>.</p>
<p>I’ve checked all the inputs, no nans or infs.</p>
<p>Anyone has any idea what’s the problem here? thank you in advance!</p>

username: ptrblck
cooked: <p>The losses should not get an NaN values, while the gradients might encounter NaNs, which will reduce the loss scaling.</p>
<p>Could you check the models operations for custom ops, which might require FP32 precision due to overflow in FP16?</p>
<p>Also, if you just want to save memory, you could try to use <code>torch.utils.checkpoint</code> to trade compute for memory.</p>

username: dhecloud
cooked: <p>Hi, thanks for your help. Yes, there was overflow in gradients for fp16 like you mentioned. I ended up using utils.checkpoint and was able to fit a larger model. Could I also trouble you one more time?</p>
<p>I ran into this error while using <code>utils.checkpoint</code>:</p>
<p><code>RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn</code></p>
<p>It happened during the <code>loss.backward()</code> function. <code>loss</code> prints as <code>tensor(0.1663, device='cuda:0)</code></p>
<p>My summarised code snippet looks like this:</p>
<pre><code class="lang-auto">outputs = checkpoint(self.t5, kwargs['input_ids'], kwargs['attention_mask'], None,\
                                 decoder_input_ids , kwargs['decoder_attention_mask'])
loss_fct = CrossEntropyLoss(ignore_index=-100)
loss = loss_fct(outputs[0].view(-1, outputs[0].size(-1)), lm_labels.view(-1))
...
loss.backward()
</code></pre>
<p>I tried following <a href="https://github.com/prigoyal/pytorch_memonger/blob/master/tutorial/Checkpointing_for_PyTorch_models.ipynb" rel="nofollow noopener">checkpoint tutorial</a> you mentioned in another thread but I couldn’t figure out what went wrong. sorry to bother you again!</p>

username: ptrblck
cooked: <p>GitHub is unfortunately down at the moment, but if I remember correctly, this issue could be solved by passing a dummy input with <code>requires_grad=True</code> to the checkpointed model?</p>
<p>PS: You could also check out the nightly binaries with <a href="https://pytorch.org/docs/master/notes/amp_examples.html#amp-examples">native amp support</a> <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"></p>

username: dhecloud
cooked: <p>Thanks! I found this <a href="https://discuss.pytorch.org/t/checkpoint-with-no-grad-requiring-inputs-problem/19117/11">link</a> and followed the same method <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"></p>