BFloat16 training - explicit cast vs autocast

username: victorx
cooked: <p>In Pytorch, there seems to be two ways to train a model in bf16 dtype.</p>
<p>One is to explicitly use input_data=input_data.to(torch.bfloat16) and model=model.to(torch.bloat16) to cast both input data and model to bfloat 16 format.</p>
<p>Another is to use torch.autocast(device_type=device, dtype=torch.bfloat16) context manager, where you don’t need to explicitly cast  the input data and model to bfloat16.</p>
<p>Are there any difference between these two options? If yes, what’s the recommended option between the two?</p>

username: ptrblck
cooked: <p>Mixed-precision training via <code>amp</code> keeps the parameters in <code>float32</code> and downcasts activations to the desired lower <code>dtype</code> for save operations. Calling <code>.to()</code> on the model directly transforms all parameters and buffers to this <code>dtype</code> and it’s the user’s responsibility to avoid under- and overflows.</p>

username: victorx
cooked: <p>Thanks <a class="mention" href="/u/ptrblck">@ptrblck</a>. I have a follow-up question about the implication of memory usage between these options.</p>
<p>According to some articles, like this one (<a href="https://huggingface.co/transformers/v4.9.2/performance.html#fp16" class="inline-onebox" rel="noopener nofollow ugc">Performance and Scalability: How To Fit a Bigger Model and Train It Faster — transformers 4.7.0 documentation</a>), training in mixed precision with fp16 actually increase the memory footprint of the model weights, as both fp32 and fp16 copies of the weights will be kept in the GPU RAM. Is this also the case for mixed precision with bf16?</p>

username: ptrblck
cooked: <p>That’s not the case for <code>float16</code> used for mixed-precision training via <code>torch.cuda.amp</code> and the copies were used in the <code>O2</code> training recipe in <code>apex</code> before <code>amp</code> was implemented directly in PyTorch.</p>

username: victorx
cooked: <p>Thanks <a class="mention" href="/u/ptrblck">@ptrblck</a> . What you said is both interesting and somewhat suprising, because even the latest version of the HuggingFace performance tuning guide (<a href="https://huggingface.co/docs/transformers/perf_train_gpu_one#fp16" class="inline-onebox" rel="noopener nofollow ugc">Methods and tools for efficient training on a single GPU</a>) still says training in fp16-based mixed precision will cause 1.5x memory usage in terms of the model weight footprint.</p>
<p>Given what you said, I was wondering whethter there is any benefit at all to do explicit-cast-to-bf16 instead of using autocast-to-bf16 <img src="https://discuss.pytorch.org/images/emoji/apple/thinking.png?v=12" title=":thinking:" class="emoji" alt=":thinking:" loading="lazy" width="20" height="20">. I have this confusion because I noticed the latest torchtune package announced by Pytorch last month used the exliclit-cast-to-bf16 option. See <a href="https://github.com/pytorch/torchtune/blob/f3611e505dfb3a7fd1bf81181ea250e63d854e0e/recipes/full_finetune_single_device.py#L246" class="inline-onebox" rel="noopener nofollow ugc">torchtune/recipes/full_finetune_single_device.py at f3611e505dfb3a7fd1bf81181ea250e63d854e0e · pytorch/torchtune · GitHub</a>, where they basically use torch.set_default_dtype(torch.bfloat16) context manager to explicitly cast data and model to bf16, and also see line 47 where they mentioned they do not use mixed precision training. Any idea why they would choose this option?</p>

username: ptrblck
cooked: <p>Explicitly casting the model to a lower <code>dtype</code> will reduce the model size as all parameters are transformed. Keep in mind that <code>amp</code> keeps all parameters in <code>float32</code> and applies casts to activations (no parameter copies are created unless you are using the deprecated <code>apex.amp</code> module). So if the training recipe and model are stable enough explicit casts can work.</p>

username: victorx
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Aha, I think I finally understand your point. To recap,</p>
<ul>
<li>old apex-based autocast to fp16/bf16 will use 1.5x memory for parameters</li>
<li>new amp-based autocast to fp16/bf16 will use 1x memory for parameters</li>
<li>explicit-cast to fp16/bf16 will use 0.5x memory for parameteters</li>
</ul>

username: lizecheng19
cooked: <p>Does this mean that when I use bfloat16 for training I can use the following code snippet to train? At the same time, I am curious whether bfloat16 does not require gradient scaling?</p>
<p>for batch in dataloader:<br>
with torch.amp.autocast(dtype=torch.bfloat16):<br>
loss = model(batch)<br>
loss.backward()<br>
optim.step()<br>
optim.zero_grad()</p>

username: ptrblck
cooked: <p>Your code is not properly formatted so check the examples in the docs and stick to their usage.<br>
Yes, <code>bfloat16</code> does not require gradient scaling.</p>

username: lizecheng19
cooked: <p>Thank you for your reply. I did not find the pytorch native bfloat16 training demo. Could you tell me how to implement it?</p>