Model training with torch_dtype=torch.bfloat16 is possible?
### System Info

- `transformers` version: 4.30.2
- Platform: Linux-5.10.184-174.730.amzn2.x86_64-x86_64-with-glibc2.26
- Python version: 3.9.16
- Huggingface_hub version: 0.15.1
- Safetensors version: 0.3.1
- PyTorch version (GPU?): 2.0.1+cu117 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No

### Who can help?

@ArthurZucker @younesbelkada 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

From #23165 and specifically this comment https://github.com/huggingface/transformers/issues/23165#issuecomment-1536439098 of @sgugger , it seems that we should not set `torch_dtype` during training. However I think this is possible. See for example the following script
```python
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
model_name = "gpt2"
input = "Hello world!"
tokenizer = AutoTokenizer.from_pretrained(model_name)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)
input_ids = tokenizer.encode(input, return_tensors="pt").to(device)
output = model(input_ids, labels=input_ids)
output.loss.backward()
optimizer.step()
print(model.get_input_embeddings().weight.grad)
```
I think gpt2 is trained using fp32 but I can load it in bfloat16 and train it (or at least get a gradient with bf16).

Thus I wonder if I misunderstood something. I have another project running to train LLaMA using bfloat16 (essentially using [`run_clm.py`](https://github.com/huggingface/transformers/blob/main/examples/pytorch/language-modeling/run_clm.py) from the official repo with `--torch_dtype=bfloat16 --bf16` command line flag), so if it turns out that I should not use `torch_dtype` for training then it means that I need to stop the experiments and rerun a lot of things :(

Thank you.

### Expected behavior

I am not sure if we should use `torch_dtype` for training.
You can try to train in full bfloat16 but it's not as stable as mixed precision bfloat16 training. When running `run_clm` with `--torch_dtype=bfloat16` you train in full bfloat16 so the flag `--bf16` (mixed precision training) is not really useful.<|||||>@sgugger thank you for the info!