Fp16 overflow when computing matmul in autocast context

username: StevenSong
cooked: <p>I’m training a gpt2 model with HuggingFace transformers and noticed I’m running into <code>nan</code> loss values during training. I tracked the source of the <code>nan</code> to a <code>softmax</code> computation where there’s a single <code>inf</code> in the input to the <code>softmax</code>. The <code>inf</code> is coming from a matrix multiply of the query and key matrices to calculate attention weights. Specifically the dot product of two vectors from query/key overflows the <code>float16</code> dtype. This is all happening within a <code>torch.cuda.amp.autocast</code> context for the forward pass of the model.</p>
<p>I’m still relatively new to <code>amp</code>, but I thought the <code>autocast</code> context was supposed to handle converting between <code>float16</code> and <code>float32</code> on the fly in cases like an overflow? If not, I wonder if some clipping would be appropriate to prevent the overflow in this specific setting, but that might be a suggestion for the HF guys <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"></p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="StevenSong" data-post="1" data-topic="183373">
<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/stevensong/48/61835_2.png" class="avatar"> StevenSong:</div>
<blockquote>
<p>but I thought the <code>autocast</code> context was supposed to handle converting between <code>float16</code> and <code>float32</code> on the fly in cases like an overflow?</p>
</blockquote>
</aside>
<p>No, <code>autocast</code> will not check for overflows, but will convert inputs/outputs to lower-precision dtypes for operations which were determined to be safe.<br>
The numerical bounds still apply and if you expect your outputs to have values outside of <code>float16</code> (<code>~65k</code>) you might need to disable it.</p>
<p>Generally, it would also be interesting to learn more about the use case and <em>why</em> such large values are expected.</p>

username: Shenyu_Lee
cooked: <p>hello, i want to know that how Pytorch amp classify the operations into the fp32 and fp16? it is firstly considered the performance on low-precision or the numerical safe?<br>
if depend on the numerical stable, the operations do the gemm also risk of overflow under fp16. And i read paper <a href="https://arxiv.org/abs/1710.03740" rel="noopener nofollow ugc">Mixed Precision Training</a> that “To maintain model accuracy, we found that some networks require that FP16 vector dot-product accumulates the partial products into an FP32 value, which is converted to FP16 before writing to memory.” It is that gemm kernel do on pytorch? thanks for your reading.</p>

username: ptrblck
cooked: <p>Internal accumulation types use <code>FP32</code> in PyTorch.<br>
Yes, values can still overflow as described in my previous post.</p>

username: Shenyu_Lee
cooked: <p>thanks you. you are nice. i need to match some operations for its appropriate datatype under mixed precision. could you give me some suggestions? how to consider about the numerical stability for operations?</p>

username: ptrblck
cooked: <p>A safe approach would be to either stick to <code>torch.amp.autocast()</code> and let it select the appropriate dtypes for the used operations or to disable <code>autocast</code> for regions you think might be sensitive to numerical stability. If you are indeed running into overflows and are also expecting such large values in your model, you could use <code>bfloat16</code> in <code>amp</code> which has a wider range than <code>float16</code> but worse precision.</p>