bf16 is more unstable than fp16, when looking at the difference of generation logprobs and forward logprobs
Author: vwxyzjn
Pull request: no

### System Info


- `transformers` version: 4.40.1
- Platform: Linux-6.5.0-28-generic-x86_64-with-glibc2.35
- Python version: 3.12.3
- Huggingface_hub version: 0.23.0
- Safetensors version: 0.4.3
- Accelerate version: 0.30.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.2.0+cu121 (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?: yes
- Using distributed or parallel set-up in script?: no

### Who can help?

@ArthurZucker

### Information

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

### Tasks

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

### Reproduction

1. Download https://gist.github.com/vwxyzjn/a9c952e0d2baf603b394f20145532d5e
2. run `python bf16_logprobs.py --fp32`
3. run `python bf16_logprobs.py --fp16`
4. run `python bf16_logprobs.py --bf16`

```
(miniforge3-22.11.1-4/envs/trl) ➜  trl git:(dataset-processor) ✗ python bf16_stuff.py --fp32
Namespace(bf16=False, fp16=False, fp32=True, seed=1)
generation_logprob[:,:5]=tensor([[-0.1527, -0.2258, -3.5535, -3.4805, -0.0519],
        [-0.1527, -0.2258, -5.2855, -0.1686, -8.4760]], device='cuda:0')
forward_logprob[:,:5]=tensor([[-0.1527, -0.2258, -3.5535, -3.4805, -0.0519],
        [-0.1527, -0.2258, -5.2855, -0.1686, -8.4760]], device='cuda:0',
       grad_fn=<SliceBackward0>)
ratio=tensor([[1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,
         1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,
         1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,
         1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,
         1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,
         1.0000, 1.0000, 1.0000, 1.0000, 1.0000],
        [1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,
         1.0001, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,
         1.0000, 1.0000, 1.0000, 1.0001, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,
         1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,
         1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,
         1.0000, 1.0000, 1.0000, 1.0000, 1.0000]], device='cuda:0',
       grad_fn=<ExpBackward0>)
ratio.mean()=1.0000026226043701
ratio.std()=1.7616826880839653e-05
ratio.max()=1.0000758171081543
ratio.min()=0.999953031539917
(miniforge3-22.11.1-4/envs/trl) ➜  trl git:(dataset-processor) ✗ python bf16_stuff.py --fp16
Namespace(bf16=False, fp16=True, fp32=False, seed=1)
generation_logprob[:,:5]=tensor([[-0.1486, -0.2212, -3.5586, -3.4688, -0.0520],
        [-0.1486, -0.2212, -5.2852, -0.1681, -8.4844]], device='cuda:0',
       dtype=torch.float16)
forward_logprob[:,:5]=tensor([[-0.1523, -0.2211, -3.5547, -3.4688, -0.0519],
        [-0.1523, -0.2211, -5.2812, -0.1667, -8.4766]], device='cuda:0',
       dtype=torch.float16, grad_fn=<SliceBackward0>)
ratio=tensor([[1.0039, 1.0000, 0.9961, 1.0000, 1.0000, 1.0059, 1.0000, 1.0000, 1.0000,
         1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,
         0.9980, 1.0000, 1.0000, 0.9961, 1.0049, 1.0000, 0.9961, 1.0000, 1.0000,
         1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,
         1.0000, 1.0000, 1.0000, 1.0000, 0.9995, 0.9990, 0.9990, 1.0000, 0.9883,
         0.9922, 0.9922, 1.0000, 0.9883, 0.9941],
        [1.0039, 1.0000, 0.9961, 0.9985, 0.9922, 1.0117, 1.0088, 1.0000, 1.0000,
         1.0059, 1.0000, 1.0000, 1.0039, 1.0000, 0.9980, 1.0000, 0.9980, 1.0000,
         1.0088, 1.0020, 0.9922, 1.0000, 0.9980, 1.0000, 1.0000, 0.9961, 1.0000,
         0.9976, 1.0078, 1.0000, 0.9990, 1.0000, 1.0000, 1.0010, 1.0098, 1.0000,
         0.9976, 0.9941, 0.9893, 1.0000, 0.9990, 1.0137, 1.0000, 0.9990, 0.9985,
         1.0000, 1.0000, 1.0127, 0.9995, 0.9980]], device='cuda:0',
       dtype=torch.float16, grad_fn=<ExpBackward0>)
ratio.mean()=1.0
ratio.std()=0.00421142578125
ratio.max()=1.013671875
ratio.min()=0.98828125
(miniforge3-22.11.1-4/envs/trl) ➜  trl git:(dataset-processor) ✗ python bf16_stuff.py --bf16
Namespace(bf16=True, fp16=False, fp32=False, seed=1)
generation_logprob[:,:5]=tensor([[-0.1426, -0.1904, -3.5938, -3.4688, -0.0618],
        [-0.1426, -0.1904, -5.2812, -0.1641, -8.5625]], device='cuda:0',
       dtype=torch.bfloat16)
forward_logprob[:,:5]=tensor([[-0.1270, -0.3457, -3.5938, -3.4688, -0.0547],
        [-0.1270, -0.3457, -5.2812, -0.1553, -8.5625]], device='cuda:0',
       dtype=torch.bfloat16, grad_fn=<SliceBackward0>)
ratio=tensor([[0.9844, 1.1719, 1.0000, 1.0000, 0.9922, 1.0000, 0.9531, 1.0938, 1.0391,
         1.0000, 1.0078, 1.1328, 1.0859, 1.0000, 0.7109, 1.0000, 1.0000, 1.0312,
         1.2422, 0.9922, 1.0312, 0.9180, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,
         1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,
         1.0000, 0.9961, 1.1328, 1.0156, 1.1172, 1.0625, 1.0000, 1.0000, 1.0000,
         1.0312, 1.0000, 0.9961, 1.0000, 1.0781],
        [0.9844, 1.1719, 1.0000, 0.9922, 1.0000, 1.0312, 0.9883, 0.9414, 1.0859,
         0.9844, 0.9922, 1.0000, 0.9102, 1.0078, 1.0000, 0.9844, 1.1016, 1.0000,
         1.0000, 1.0000, 1.0000, 0.9688, 0.9531, 1.0078, 1.0000, 1.0000, 0.9844,
         1.0156, 0.9766, 0.9844, 0.9961, 1.0000, 0.9961, 1.0000, 1.0000, 1.0000,
         1.0078, 1.0000, 1.0000, 0.9531, 1.0000, 1.0000, 1.0078, 1.0000, 1.0000,
         0.9688, 1.0078, 1.0000, 1.0000, 1.0000]], device='cuda:0',
       dtype=torch.bfloat16, grad_fn=<ExpBackward0>)
ratio.mean()=1.0078125
ratio.std()=0.05712890625
ratio.max()=1.2421875
ratio.min()=0.7109375
```

### Expected behavior

Basically, the ratio should be 1s

Comments:
I also encounter more instability with bf16 in rope related models like pythia (i.e. the gptneox architecture). This occurred to me when working on SDPA support for the architecture at #31031.

My guess it is due to rope which forces upcasting to fp32. Further down the line it is usually downcasted somewhere along there which is more severe for bf16 (as it has less precision) which results into this phenomenon. For example, this upcasting is also circumvented in the flash attention 2 variant by downcasting:
https://github.com/huggingface/transformers/blob/940fde8dafaecb8f17b588c5078291f1c1a420c8/src/transformers/models/gpt_neox/modeling_gpt_neox.py#L364-L369

Edit: tl;dr: more precision is needed which bf16 can't offer

Yep, I think you should check if #29285 applied to the model you are using does not solved the issue? 🤗 

I tried applying the same fix in #29285 naively but it doesn't seem to make a difference

<img width="1752" alt="image" src="https://github.com/huggingface/transformers/assets/5555347/f6fc74c8-b905-4531-84cd-214d10b6bb5e">


Did some further digging, the logprobs of the forward pass feel length related. For example, if I run the model through the first 20 generated tokens, I get 

```
(.venv) ➜  transformers git:(main) ✗ python -i bf16_stuff.py --bf16
Namespace(bf16=True, fp16=False, fp32=False, seed=1)
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
A decoder-only architecture is being used, but right-padding was detected! For correct generation results, please set `padding_side='left'` when initializing the tokenizer.
generation_logprob[:,:5]=tensor([[-0.1426, -0.1904, -3.5938, -3.4688, -0.0618],
        [-0.1426, -0.1904, -5.2812, -0.1641, -8.5625]], device='cuda:0',
       dtype=torch.bfloat16)
forward_logprob[:,:5]=tensor([[-0.1426, -0.1904, -3.5781, -3.4062, -0.0618],
        [-0.1426, -0.1904, -5.2812, -0.1670, -8.5625]], device='cuda:0',
       dtype=torch.bfloat16, grad_fn=<SliceBackward0>)
ratio=tensor([[1.0000, 1.0000, 0.9844, 0.9375, 1.0000, 0.8945, 1.0000, 1.0000, 1.0000,
         1.0000, 1.0000, 1.0000, 1.0859, 1.0000, 1.0000, 1.0000, 1.0000, 1.0312,
         1.2266],
        [1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 0.9961, 1.0859,
         1.0625, 0.9922, 1.0000, 0.9688, 1.0000, 1.0000, 1.0000, 0.9844, 1.0000,
         1.0000]], device='cuda:0', dtype=torch.bfloat16,
       grad_fn=<ExpBackward0>)
ratio.mean()=1.0078125
ratio.std()=0.048095703125
ratio.max()=1.2265625
ratio.min()=0.89453125
```

If I do the forward pass on the first 40 generated tokens, I get

```
generation_logprob[:,:5]=tensor([[-0.1426, -0.1904, -3.5938, -3.4688, -0.0618],
        [-0.1426, -0.1904, -5.2812, -0.1641, -8.5625]], device='cuda:0',
       dtype=torch.bfloat16)
forward_logprob[:,:5]=tensor([[-0.1270, -0.3164, -3.5938, -3.3906, -0.0547],
        [-0.1270, -0.3164, -5.1562, -0.1387, -8.5625]], device='cuda:0',
       dtype=torch.bfloat16, grad_fn=<SliceBackward0>)
ratio=tensor([[0.9844, 1.1328, 1.0000, 0.9258, 0.9922, 0.9688, 0.9023, 1.0000, 1.0000,
         1.0000, 0.9336, 1.0000, 0.8477, 1.0000, 0.8047, 1.0000, 1.0000, 0.9961,
         1.2266, 1.0078, 0.9688, 0.9258, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,
         1.0000, 1.0234, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 0.9961,
         1.0000, 1.0000, 1.0000],
        [0.9844, 1.1328, 0.8828, 0.9766, 1.0000, 1.0000, 0.9727, 0.9375, 0.9922,
         1.0000, 1.0000, 1.0000, 0.8828, 1.0078, 1.0000, 0.8672, 1.0156, 1.0000,
         1.0000, 1.0000, 1.0156, 1.0000, 0.8164, 1.0156, 1.0078, 1.0000, 0.9844,
         1.0234, 1.0000, 0.9688, 0.9961, 1.0000, 0.9961, 1.0234, 1.0000, 1.0000,
         1.0000, 1.0000, 1.0000]], device='cuda:0', dtype=torch.bfloat16,
       grad_fn=<ExpBackward0>)
ratio.mean()=0.98828125
ratio.std()=0.057373046875
ratio.max()=1.2265625
ratio.min()=0.8046875
```

Notice how the logprob changed from `-0.1426` to `-0.1270`.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Hi, is there any updates? Thanks!

This is basically fixed in most models (the rope part for example) the rest can be traced back to https://github.com/huggingface/transformers/issues/25420#issuecomment-1775317535