Bfloat16 training

nastya236: <p>Hey!<br>
I was wondering about your experience with bf16 training, specifically the type of the gradients and the weights.<br>
From the repositories I’ve checked, it seems there are two main approaches (assuming a single GPU, no FSDP, and all activations are recomputed during backward):</p>
<ol>
<li>
<p>BF16 weights + BF16 grads + FP32 master copy + FP32 optimizer state → (2 + 2 + 4 + 8) = 16 bytes/param (as in <a href="https://docs.nvidia.com/nemo-framework/user-guide/24.09/nemotoolkit/features/mixed_precision.html" rel="noopener nofollow ugc">NeMo</a>).</p>
</li>
<li>
<p>BF16 weights + FP32 grad accumulation + FP32 optimizer state→ (2 + 4 + 8) = 14 bytes/param (I haven’t found a reference, but it appears Megatrom-LM does this).</p>
</li>
</ol>
<p>I then questioned the need for FP32 grad accumulation / copy of the weights so using just bf16 weights, bf16 gradients, optimizer 32 = 12 bytes and came across <a href="https://arxiv.org/pdf/2405.03637" rel="noopener nofollow ugc">this paper</a> (page 6, Section 5.1 “Precision options”), case D−MW - bf16 weights, bf16 gradients, optimizer 32 no master copy. But according to Table 3, performance is worse compared to the first option.<br>
So I was wondering if you experienced any performance degradation if gradients, weights are in bf16 only? Why it is related to pytorch forum:</p>
<p>The question:</p>
<p>should I do:<br>
<code>model.to(bfloat16); x.to(bfloat16)</code> – so no master copy, gradients are bf16 type<br>
or<br>
<code>with torch.autocast('cuda', dtype=torch.bfloat16):</code><br>
<code>   losses = loss_fn(model, x, y)</code>  – still have master copy</p>
<p>Thank you very much for your input!</p>

J_Johnson: <p>Best to stick with autocast.</p>
<ol>
<li>The 2nd configuration you mentioned is not 14-bytes/param. Should be 16-bytes/param: BF16 compute weights(2 bytes) + FP32 master weights(4 bytes) + BF16 gradients(2 bytes) + FP32 optimizer states(8 bytes).</li>
<li>As the paper you linked demonstrates, even with FP32 optimizer states, skipping a master copy degrades effective descent quality.</li>
</ol>

Arunprakash-A: <aside class="quote no-group" data-username="nastya236" data-post="1" data-topic="222382">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/nastya236/48/77968_2.png" class="avatar"> nastya236:</div>
<blockquote>
<p>So I was wondering if you experienced any performance degradation if gradients, weights are in bf16 only?</p>
</blockquote>
</aside>
<p>Yes. It will hurt performance, as a significant number of weight values become zero when cast to bf16 (fp16), as shown in the figure below.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/5/7551e522b91ce925fe1156ceeccaa9576e582dc6.png" data-download-href="https://discuss.pytorch.org/uploads/default/7551e522b91ce925fe1156ceeccaa9576e582dc6" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/7/5/7551e522b91ce925fe1156ceeccaa9576e582dc6.png" alt="image" data-base62-sha1="gJRk39G9OCvu8AnV4fXpJfZYZVQ" width="511" height="397"><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">511×397 14.1 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 are various design choices to overcome this issue. For more information, you may refer to<a href="https://nvidia.github.io/apex/amp.html" rel="noopener nofollow ugc">this page</a> on NVIDIA’s Amp. As suggested by <a class="mention" href="/u/j_johnson">@J_Johnson</a>, it is preferable to use <code>autocast</code>.</p>