Huggingface microsoft/mdeberta model never gets updated under AMP

Fusionastic: <p>I’m using the Huggingface <a href="https://huggingface.co/microsoft/mdeberta-v3-base" rel="noopener nofollow ugc">microsoft/mdeberta-v3-base</a> pretrained model to finetune on my task. I discovered that during training, if using AMP, the optimizer never gets updated. step() method of the optimizer never gets called by scaler, because the scale of the scaler keeps shrinking until this float point number itself gets underflowed. So the model cannot be updated. Once turned off AMP, everything works fine.</p>
<p>Moreover, I tried finetuning other pretrained models from huggingface, like    <a href="https://huggingface.co/microsoft/deberta-v3-base" rel="noopener nofollow ugc">microsoft/deberta-v3-base</a> and more, with AMP on, and they all work fine. The scale would typically hold at 16384.</p>
<p>I’m pretty sure that my code for AMP usage has nothing problematic since it only fails on this particular pretrained model. I tried increasing the init_scale of GradScaler, from the default <code>2**16</code> to <code>2**40</code> with exponent interval 4, and they all failed.</p>
<p>I wonder if there is a workaround for this. This particular pretrained model, mdeberta-v3-base, is pretty strong, and we don’t wanna give it up, nor training it with AMP off, which costs a huge amount of time.</p>

ptrblck: <aside class="quote no-group" data-username="Fusionastic" data-post="1" data-topic="155162">
<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/fusionastic/48/50908_2.png" class="avatar"> Fusionastic:</div>
<blockquote>
<p>the optimizer never gets updated. step() method of the optimizer never gets called by scaler, because the scale of the scaler keeps shrinking until this float point number itself gets underflowed. So the model cannot be updated. Once turned off AMP, everything works fine.</p>
</blockquote>
</aside>
<p>This could mean that the forward pass or the loss calculation would create invalid values and the <code>GradScaler</code> tries to decrease the scaling factor in case it’s too large.<br>
However, if the output is already invalid, decreasing the scaling factor would of course not work anymore, so you shoulde definitely check is the predictions or loss already contains Infs or NaNs.</p>
<p>If so, check if the model was pretrained with <code>bfloat16</code> as it has a larger range and less precision than <code>float16</code>. In that case, you might want to use <code>autocast</code> with <code>dtype=torch.bfloat16</code>.</p>
<aside class="quote no-group" data-username="Fusionastic" data-post="1" data-topic="155162">
<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/fusionastic/48/50908_2.png" class="avatar"> Fusionastic:</div>
<blockquote>
<p>I tried increasing the init_scale of GradScaler, from the default <code>2**16</code> to <code>2**40</code> with exponent interval 4, and they all failed.</p>
</blockquote>
</aside>
<p>This wouldn’t help and would create overflows more easily.</p>

Fusionastic: <p>Thanks for your informative advice. The loss is normal, so maybe there are some NaN or infs in gradient or params. Using torch.bfloat16 solves the problem, though it comes alongside with some performance decrease. Besides, not all GPUs support bfloat16, including V100.</p>