That said, also be aware that if you pre-trained a model in bf16, it's likely to have overflow issues if someone tries to finetune it in fp16 down the road. So once started on the bf16-mode path it's best to remain on it and not switch to fp16.
In 🤗 Transformers bf16 mixed precision is enabled by passing --bf16 to the 🤗 Trainer.
If you use your own trainer, this is just:
from torch.cuda.amp import autocast
with autocast(dtype=torch.bfloat16):
    loss, outputs = ...
If you need to switch a tensor to bf16, it's just: t.to(dtype=torch.bfloat16)
Here is how you can check if your setup supports bf16:
python -c 'import transformers; print(f"BF16 support is {transformers.file_utils.is_torch_bf16_available()}")'
On the other hand bf16 has a much worse precision than fp16, so there are certain situations where you'd still want to use fp16 and not bf16.
