Fix llava half precision and autocast issues
Author: frasermince
Pull request: yes
Labels: name: run-slow

What does this PR do?

Currently half precision training on LLaVa is broken due to the language embedding output always being full precision. This PR matches the behavior of the original LLaVa implementation by casting the result of language embedding to the image_features' dtype. See [here](https://github.com/haotian-liu/LLaVA/blob/main/llava/model/multimodal_encoder/clip_encoder.py#L51) for relevant line in the original repo.


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker and @younesbelkada
- vision models: @amyeroberts
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @gante
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @pacman100

Integrations:

- deepspeed: HF Trainer/Accelerate: @pacman100
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc and @younesbelkada

Documentation: @stevhliu and @MKhalusova

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
> Thanks for adding this!
> 
> Could you add a test e.g. like `create_and_check_reformer_model_fp16_forward` that we have for other models?

Apologies on the wait on this. After implementing the test, it looks like the problem here is not quite what I thought it was. I will try to respond soon with either fixing the correct problem or closing this PR.

After some further investigation, I tracked down exactly what was happening. `.half()` or `.to(torch.bfloat16)` work fine on the model. However, the issue arises when the model is run within `with torch.autocast`. This also includes if you are using a trainer with the `bf16` or `fp16` arg set to true since `accelerate` uses an `autocast` context behind the scenes. When run in autocast bf16 or fp16 block the [language model embedding](https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava/modeling_llava.py#L419) has full precision. However, the [`image_features`](https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava/modeling_llava.py#L436) are in half precision. Due to these being merged there is a type error. To fix this we need to either cast embeddings to half precision or image features to full. It made the most sense to me to cast the embeddings to half precision. Let me know if you prefer we go the other direction but this seemed the most efficient.

I also fixed the same issue on llava_next. In the process I discovered that there were NaNs occuring upon running llava_next in half precision. To fix this I used `randn` instead of `empty` similarly to the [original llava initialization](https://github.com/haotian-liu/LLaVA/blob/4e2277a060da264c4f21b364c867cc622c945874/llava/model/llava_arch.py#L85). 


I’ll need to dig into this when I’m back to my computer, but @amyeroberts just to describe how the autocast system works in accelerate quickly:

We wrap around the models’ .forward() with an autocast manager, so that the original weights remain in fp32. Not doing so leads to a large deal of precision and end accuracy issues, which is why we maintain this reference. It uses a pinch more memory (than pure doing `model.half()`) but for good reason. And then we convert the output back to fp32 at the end

If certain weights need to be cast for certain models, I think that can be okay and taken as a case-by-case basis, but I want to re-read this discussion carefully to understand why. 

(And we should test training for a few epochs/steps to make sure accuracies etc remain as expected when using pretrained weights)

> I’ll need to dig into this when I’m back to my computer, but @amyeroberts just to describe how the autocast system works in accelerate quickly:
> 
> We wrap around the models’ .forward() with an autocast manager, so that the original weights remain in fp32. Not doing so leads to a large deal of precision and end accuracy issues, which is why we maintain this reference. It uses a pinch more memory (than pure doing `model.half()`) but for good reason. And then we convert the output back to fp32 at the end
> 
> If certain weights need to be cast for certain models, I think that can be okay and taken as a case-by-case basis, but I want to re-read this discussion carefully to understand why.
> 
> (And we should test training for a few epochs/steps to make sure accuracies etc remain as expected when using pretrained weights)

Sounds good! Let me know if I can do anything to make the reason for these changes more clear!

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

> Overall I think these are good tests how they are, let's please keep a close eye on it and see what other models have issues. We may want to change some things with Accelerate (like specifying certain layers to be in the precision fully) if tests show no loss in performance :)

Sounds good! I will also add the trainer test in a separate PR as requested.

@muellerzr before merging do you have any suggestions on ensuring accuracies remain as expected? Is it basically taking the assertions in some of the integration tests after a couple of training steps?

>  Is it basically taking the assertions in some of the integration tests after a couple of training steps?

Yes indeed! (Assertion or `torch.close` more than likely)

> ```
> RUN_SLOW=1 pytest tests/models/llava tests/models/llava_next
> ```

There was one that seemed to be already failing before my changes. I took the liberty of changing the assertion slightly. Let me know if this is acceptable! There are a couple failing in main for me for llava_next as well but that would require me changing the expected logits. Not sure if that is acceptable or not.

@amyeroberts leave to your discretion :) 

Given slow the test was failing on main I changed the logits in the test for now. If that is wrong I can remove.

@frasermince Great - thanks for checking! I'm going to be a bit annoying here - as the differences are so small it's quite possible they're just coming from running on different hardware / slightly different env setups. So we might need to revert. 

@ydshieh How can we run the slow tests on this branch in the new github actions workflow to make sure the values are aligned? 

@amyeroberts I am on it and should be ready soon (even possible today 🔥 )

It is in #30540 :-)

@frasermince Thanks for your patience! 

We've just merged in #30540 (thanks @ydshieh!) so should be able to run the slow tests on our runners. To trigger this could you: 
* Rebase on `main` to include the recent commit 
* Push an empty commit with the message `[run-slow]llava,llava_next`

Thanks for triggering the tests! OK, so we're hitting OOM errors, which is going to be a pain to try and resolve as it's a wider issue. What I suggest is:
* We resolve the conflicts
* We merge this in 
* If there's some issues with the slow tests which aren't oom, I'll handle in a follow up PR (cc @ydshieh - I'll be the person to ping)

@frasermince Thanks for all of your patience iterating with us. We're good to merge!