RoPE loses precision for Llama / Gemma + Gemma logits.float()
Author: danielhanchen
Pull request: yes

Tagging @ArthurZucker :)

When I was implementing Gemma for Unsloth, I noticed when one uses bfloat16, the RoPE embeddings get autocast to bfloat16, when we require it to be in float32. This causes the positional encodings to lose precision dramatically especially for very large context lengths.

Below I pasted the image on how HF for now handles RoPE. You can see the loss in precision when using bfloat16. I manually autocasted it to float32 in Unsloth, and you can see the expected positional encodings.
![image](https://github.com/huggingface/transformers/assets/23090290/ce691e66-ac77-4904-a4f1-ecf15db48a16)

I couldn't find why Unsloth's error could not match that of HF's original Gemma implementation. On float16, this issue does not occur, with HF and Unsloth's training loss curve being equivalent:
![image](https://github.com/huggingface/transformers/assets/23090290/eda26d8a-2102-4f17-baaa-dc6b9d93d799)

However when I switched over to bfloat16, HF and Unsloth's training losses diverge at the start, and Unsloth always retains a lower loss as training goes on:
![image](https://github.com/huggingface/transformers/assets/23090290/179c0d31-45f9-4109-8eb1-33ed0a7b74b9)

If you look at the losses more carefully (same seed), you can see the differences more closely.
![image](https://github.com/huggingface/transformers/assets/23090290/22b94ff9-96ac-4d15-9c18-ede35e3920c8)

The culprit I found was
```python
inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
position_ids_expanded = position_ids[:, None, :].float()
freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
emb = torch.cat((freqs, freqs), dim=-1)
```
where if one uses `torch.autocast()`, `freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)` gets done in bfloat16 and not float32. I propose we turn off autocast to force float32. Ie:
```python
with torch.autocast(device_type=position_ids_expanded.device.type, enabled=False):
    freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
```
This ensures `torch.autocast` to turn off automatic downcasting to float16 / bfloat16 for the RoPE embeddings. My proposed fix shows the following loss curve:
![image](https://github.com/huggingface/transformers/assets/23090290/fefe0663-ec79-4cac-8ffe-dfc932bd09f3)

Also, in Gemma, a 1 liner was missed :) `logits = logits.float()` must be placed to upcast the `logits` to float32. Although it should be done automatically in `torch.autocast`, it's best to keep the convention as done in llama, mistral and other models. Gemma's implementation seems to maybe have forgotten this 1 line :)

Comments:
Forgot to add I'm not certain if this will break CUDAGraphs for faster inference - hopefully not

Sadly unsure if it's just for training :(( For inference I don't remember up to which context length, bfloat16 won't be an issue. I think it was up to 4096. However, bfloat16 loses precision even for inference sadly after 4096 context lengths. 8192 definitely - bfloat16 essentially thinks the last 4 tokens are all position 8192 ie [8192, 8192, 8192, 8192], whilst the correct float32 is [8188, 8189, 8190, 8191].

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_29285). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@danielhanchen `.sin()` and `.cos()` should ideally happen in FP32 as well. Have you noticed any performance changes if you force them to happen in FP32?

@gante Actually interesting point - I can see torch.autocast does arc sin and sinh etc in float32, but it doesnt list sin itself - I'll have to check if .sin() is done in float32 or float16

@ArthurZucker I checked everything and it's working! You guys can double check if anything is wrong. You can push the commit whenever. Thank you! :) 

Ugh that 1 CI issue

@danielhanchen we may need a conditional context manager 😢 

cc @michaelbenayoun 

Alright the failing test is now related to cos and sin so fixed the device thing at least

```python
        self.register_buffer("_cos_cached", emb.cos().to(torch.get_default_dtype()), persistent=False)
        self.register_buffer("_sin_cached", emb.sin().to(torch.get_default_dtype()), persistent=False)
```
 I doubt that this is the cause of the issue since we used to do that before alrady

Thanks @danielhanchen 🤗 merging now!

Yay! :) Great work everyone :)

(hello from the Gemma team!) Superb stuff, thanks for this lovely fix :) 

@suryabhupa Thanks :)
Actually was gonna ask do you know if Gemma uses approximate gelu or exact gelu? When comparing Keras to HF, torch.dist gets 4.7943, while tanh approx gets 0.0057

![image](https://github.com/huggingface/transformers/assets/23090290/159f6d78-c422-4885-84a8-c3853b27b21e)


Did anyone realize that this fix would incur additional VRAM increase likely making it OOM for the code that could have been trained?

@paulcx Yes unfortunately that was a consideration we took - the issue is results become incorrect especially on longer context lengths - shorter is fine. We tried our best to isolate changes - for inference this is fine, but for training there will be more VRAM usage. There's always Unsloth for Gemma of course, where we allow 2.5x faster finetuning and 60% VRAM reductions :) And TRL, PEFT all HF libraries are all integrated. Gemma 7b notebook: https://colab.research.google.com/drive/10NbwlsRChbma1v55m8LAPYG15uQv6HLo?usp=sharing

And though this increases VRAM usage, precision and quality of the outputs should be improved. This is also more "accurate" with respect to the original implementations

@danielhanchen Thank you for your finding. It would be exciting if it could be modified on top of the existing one without at least increasing the VRAM. At least, some resource-strained tranning may not be able to train now due to this VRAM (even if the batch size is set to 1).

How much of a VRAM increase are we talking about? 

It's difficult to give an exact number, but the growth in VRAM requirements could depend on the size of the model, for instance, for a 34B model, it could be in the tens of GBs？

@paulcx Wait 10s of GBs? The PR is only for Llama and Gemma, so I'm assuming CodeLlama 34b is using more VRAM? The RoPE upcasting should only use 8192 * 8192 * 16 bits approx of extra VRAM * n layers so say 32 layers = 128MB of extra VRAM per layer, and it should be cleared away since we don't need the 8192x8192 matrix. I don't see how an extra 10GB of VRAM is coming from - are you finetuning Gemma or Llama?

I can see why Gemma might be using more VRAM - since `logits.float()` is necessary otherwise the softmax will use `torch.float16` which causes incorrect results when training over long periods of time.

For Llama, only at max 128MB of VRAM should be used extra

@danielhanchen I must say that my estimate is likely to be quite imprecise given that I have not made accurate experiments and statistics. The only evidence I have to support GBs is from my own experience with a 34B CodeLlama  (which is trained successfully on transformers==4.37.2). I know it could be caused by many reasons, for example, multi-GPU training with deepspeed ... etc.

![309605687-962a3bf3-db1f-4c10-b478-1de18ddbafb6](https://github.com/huggingface/transformers/assets/738834/eddd04e3-e6c1-42b3-848d-2ca8d626771c)


@paulcx OHH this is a different issue!! It seems like a change was made in another PR which allocates a causal mask of size (16384, 16384) https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L940

The triu causes the causal mask to upcast to float32, using 16384^2 * 4bytes = 1GB of extra VRAM. We have `n^2 * 4 / 1024 / 1024 = 37.25GB` in your screenshot, so I'm assuming you're also doing RoPE Scaling to 100K context length? So ie a (100K, 100K) matrix was trying to be created.

So this looks like a separate problem than this PR unfortunately.

Thanks @danielhanchen 

@paulcx your issue is related to this one (https://github.com/huggingface/transformers/issues/29484) -- let's keep the discussion there! :)

FYI -- as reported in https://github.com/pytorch/pytorch/issues/128394, this change seems to break torch.export of the Llama 2 model. There seems to be two causes of the break:

1. `torch.autocast` does not support `meta` device, as suggested in its documentation:
> device_type ([str](https://docs.python.org/3/library/stdtypes.html#str), required) – Device type to use. Possible values are: ‘cuda’, ‘cpu’, ‘xpu’ and ‘hpu’. The type is the same as the type attribute of a [torch.device](https://pytorch.org/docs/stable/tensor_attributes.html#torch.device). Thus, you may obtain the device type of a tensor using Tensor.device.type.

Creating a model in `meta` mode seems to be quite useful for quite a few applications, in particular when materializing the whole model on a single device is infeasible.

2. `torch.export` today does not work well with `torch.autocast`, as `torch.amp.autocast_mode._enter_autocast` is not a valid ATen op. See detailed [discussion](https://github.com/pytorch/pytorch/issues/128394#issuecomment-2221075748) here.

Wondering if there is a way to reenable export of Llama 2? Thanks!

Cc: @angelayi @lessw2020

Hi @kwen2501 👋 Happy to iterate with you to get a working solution for `torch.export` 🤗 

I'm not experienced with `torch.autocast`, so I have a question for you. Our current code has

```py
        with torch.autocast(device_type=device_type, enabled=False):
            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
```

Isn't `.float()` and the context manager disabling autocast redundant for autocasting purposes? If so, given that `freqs` would always be a `float32` tensor, removing the context manager should do the trick

Hi @gante and @kwen2501 and @angelayi,
@gante  - you are correct that the autocast here is redundant both since it is disabled and the relevant variables are force upcast to float anyway, further making the autocast redundant.

(to your question, short summary is autocast will dynamically lower the precision to fp16 where it feels it can do so without losing precision. In this case it appears that for long context the lower precision could be negative and hence the reason to ensure the freqs calc is always done in fp32).  

To try and speed things up in resolving, I've made a PR that removes the device type calcs that were being fed to the disabled autocast and of course removed the autocast itself.


When I tested this on Friday, removing the autocast did enable torch.export to function as expected so this should resolve this for us. 

But appreciate any feedback on the PR in case I have missed something! 
(adding @danielhanchen to the PR reviewers as well since he isolated the original core issue)



Note - PR is failing CI due to this - @gante, would you be able to help review this error?
~~~
File "/root/transformers/utils/check_copies.py", line 856, in check_copies
    raise Exception(
Exception: Found the following copy inconsistencies:
- src/transformers/models/mistral/modeling_mistral.py: copy does not match models.llama.modeling_llama.LlamaRotaryEmbedding.forward at line 95
- src/transformers/models/olmo/modeling_olmo.py: copy does not match models.llama.modeling_llama.LlamaRotaryEmbedding at line 97
Run `make fix-copies` or `python utils/check_copies.py --fix_and_overwrite` to fix them.
~~~

I'll comment here as well - the primary reason why the context manager was placed was because mixed precision training would cause `torch.amp.autocast` to randomly ignore the `.float()` upcast, and so the solution was to force the autocast to get disabled in that region - if newer Torch versions don't downcast, then the autocast can be safely removed.

Thanks @gante @lessw2020 @danielhanchen !
It seems the autocast disabler was intended to work in a bigger autocast enabling region (added by a user).

I wonder if the user can divide the autocast enabling region into two to skip this disabled region? Perhaps we can a note to the `LlamaRotaryEmbedding` module to warn that applying autocast would cause precision loss?

Longer term, maybe it would be worth adding a `torch.no_autocast` API -- the current way of `enabled=False` is a bit confusing. 

Thanks @danielhanchen for the insight to the reasoning behind the autocast w/ enabled = False, here (I suspected there was likely more to the story).

I did want to add that we had the inverse of this situation with fsdp - for fsdp mixed precision, if you explicitly set a model dtype to something esp .bfloat16, and then pass to FSDP, we respect that.  That behaviour however, then generated a big commotion from HF about why wouldn't we secretly auto override this for mixed precision and keep fp32 weights. (because we respect what the user has directed).
Anyway, it sems the behaviour is 180 degree opposite here - if a user sets .float(), then that should be respected as a direct dtype command...not just something to be arbitrarily overridden without a blocking disabled autocast.  
Regardless, agree with @kwen2501  that there needs to be a direct way for the user to spec a dtype and not have that secretly overridden, which would then be consistent with mixed precision in FSDP (i.e. you spec it, you got it, the end). 


>that there needs to be a direct way for the user to spec a dtype and not have that secretly overridden, which would then be consistent with mixed precision in FSDP (i.e. you spec it, you got it, the end).

do you mean for ROPE? 

https://github.com/huggingface/transformers/pull/31959#issuecomment-2230537392

"What I'm reading then is that we SHOULDN'T merge this PR until we confirm that pytorch doesn't change the type of an explicit .float() cast when autocast is active. Otherwise, we will get a regression (thank you @danielhanchen for confirming 🤗 )"

Hi all - we have a much better solution now by making some changes to how PP does the model tracing via PT export and with that, it now handles the autocast issue directly so no changes needed here in the transformer code.
Thus, I'm going to close my earlier PR as the proper fix is within export/PP directly.
The PR for reference is here in PyTorch:
https://github.com/pytorch/pytorch/pull/130998

Thanks @danielhanchen @gante @ArthurZucker @kwen2501  for the details and convo on this.  
I'm glad we are able to have a more fundamental fix for this! 

Thanks for updating, great to know that overall this will be seamlessly fixed for everyone! 🤗