RoPE loses precision for Llama / Gemma + Gemma logits.float()

[danielhanchen]

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.

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:

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:

If you look at the losses more carefully (same seed), you can see the differences more closely.

The culprit I found was

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:

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:

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 :)

[danielhanchen]

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

[ArthurZucker]

I'll have to check the compile test and everything, but we usually hate these kind of changes 🫣 the bug is real, I'll see if I can find a good alternative as this is pretty much only for training! Great catch 🤗

[danielhanchen]

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].

[ArthurZucker]

LGTM let's do no grad and autocast, I'll test compile once you have both!

[ArthurZucker]

LGTM, before merging I'll ping @pacman100, @younesbelkada and @fxmarty as this is pretty important! Feel free to comment if you are against these changes!

[ArthurZucker]

Suggested change

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[gante]

@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?

[danielhanchen]

@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

[danielhanchen]

@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! :)

[ArthurZucker]

LGTM! Thanks a mile for this.
Let's make sure you run make style and make fixup for the last CIs

[kwen2501]

FYI -- as reported in pytorch/pytorch#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, 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. 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 here.

Wondering if there is a way to reenable export of Llama 2? Thanks!

Cc: @angelayi @lessw2020

[gante]

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

        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

[lessw2020]

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)

[lessw2020]

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.

[danielhanchen]

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.

[kwen2501]

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.

[lessw2020]

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).

[ArthurZucker]

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?

[gante]

#31959 (comment)

"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 🤗 )"

[lessw2020]

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:
pytorch/pytorch#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!

[ArthurZucker]

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