[core] enable bf16 training

[younesbelkada]

Currently on the main branch of trl training in bfloat16 & logging with wandb fails. Users will encounter an error that are hard to interpret.

The fixes are the following:

• remove wandb.watch(model) as it takes care of logging gradients that are in bf16, otherwise users will get an error that is similar to [CLI]: RuntimeError: "histogram_cpu" not implemented for 'BFloat16' wandb/wandb#3332 | I propose to remove this is not a core feature and is a main blocking point
• upcast the tensors in fp32 before logging them
• force-upcast the lm_logits in fp32

cc @lvwerra

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[lvwerra]

What's the issue with lm_logits being in bf16?

[younesbelkada]

So that we can compute the loss in fp32 and I found that this was more stable + sometimes we log a list of tensors, and directly casting the loss in fp32 avoids the issue with numpy & bf16

[RylanSchaeffer]

@younesbelkada @lvwerra I'm running into a problem where this forced upcast causes a massive spike in memory for models with large vocabularies (e.g., Gemma 2 by Google). This then either throws an OOM error or forces me to cut the minibatch size in half, which doubles the PPO runtime

Issue: #1980

Could you please provide more information or evidence about the stability argument?