[RFC] Support full bf16 training

[ebsmothers]

This PR adds support for full bfloat16 training. In SFT it is pretty common to store everything in bfloat16 to save memory, with select tensors (logits, RoPE buffers and activations) maintained in a higher precision to preserve numerical accuracy. Separately I think having this supported more generally would be useful for faster iteration -- e.g. it allows me to run Llama3 70B on a single node of H100s, which otherwise is not possible with the default config.

Assuming this is generally useful, would like feedback on:

1. Acceptable loss convergence: in the first 100 steps on Llama3 8B full bf16 training goes from 12.25 -> 8, as opposed to 12.25 -> 7 with fp32 training. Is this a concern? (As mentioned, for SFT this is less of an issue; happy to validate that statement if that's helpful.)
2. Interaction with mixed precision training -- where is the right place to validate that these are not both set at once?
3. Where to put the set_default_dtype API

[fegin]

If we set the training dtype during the training initialization, why not also do the output conversion in the trainer (train loop)?

[tianyu-l]

This is already in the loss function https://github.com/pytorch/torchtitan/blob/main/torchtitan/components/loss.py#L21

Also see #642

[ebsmothers]

Thanks, just removed

[fegin]

What if mixed_precision_param is float32 but dtype is bfloat16? There should be a check?

[ebsmothers]

Yeah agreed. Do we want to do this somewhere in train.py? Lmk if you think there's a better place

[tianyu-l]

mixed_precision_param is coming from FSDP2. I think if FSDP2 can work with that, it's users responsibility to config them properly.

[fegin]

We also make it work with DDP/single device: #1303. I think a warning is at least required.

[ebsmothers]

Sounds good. In that case I will leave this as is

[tianyu-l]

@fegin
autocast is not well supported in torchtitan anyways. I'm not sure if it is still maintained. See other issue like #1525

But sure, having a warning sounds good.