When using 4.35 to load llama2 and flash attention2 for training, the memory usage has nearly quadrupled
Author: Trangle
Pull request: no

### System Info

transformers version: 4.35.0
Platform: Ubuntu22.04
Python version: 3.11
Huggingface_hub version: 0.17.3
Safetensors version: 0.4.0
Accelerate version: 0.24.1
Accelerate config: not found
PyTorch version (GPU?): 2.1.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?

When switched to 4.34, the memory usage decreases by nearly four times

### Information

- [ ] The official example scripts
- [ ] 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

Some example like: https://github.com/lm-sys/FastChat/blob/main/scripts/train_vicuna_7b.sh

or any training process using 4.35

### Expected behavior

Code does not run out of CUDA memory.

Comments:
Hi @Trangle, thanks for raising this issue! 

For the diff highlighted in the issue, this shouldn't affect the memory consumption of the model. This is just a change in type annotation. 

Could you provide some more details about the setup and what's being observed? Specifically: 
* What hardware are you using?
* Is 4x memory consumption the limit of the GPU i.e. are you hitting OOM before seeing the max load?  
* At what line of code is this explosion in memory being observed? e.g. when the model's loading? On the forward pass? 

We have many requests for help, which we can only attend at a decent pace if you help us too. The linked to script is large and has a lot of logic. Could you provide a code snippet we can run to reproduce the error with the minimal amount of code? 

> Hi @Trangle, thanks for raising this issue!
> 
> For the diff highlighted in the issue, this shouldn't affect the memory consumption of the model. This is just a change in type annotation.
> 
> Could you provide some more details about the setup and what's being observed? Specifically:
> 
> * What hardware are you using?
> * Is 4x memory consumption the limit of the GPU i.e. are you hitting OOM before seeing the max load?
> * At what line of code is this explosion in memory being observed? e.g. when the model's loading? On the forward pass?
> 
> We have many requests for help, which we can only attend at a decent pace if you help us too. The linked to script is large and has a lot of logic. Could you provide a code snippet we can run to reproduce the error with the minimal amount of code?

I use 8 * A100-80G, with seq_length=4096，flash-attention2

I'm not sure which part caused the surge in graphics memory and caused OOM, but this issue was resolved by downgrading 4.34.1

> I use 8 * A100-80G, with seq_length=4096，flash-attention2

Thanks!

> I'm not sure which part caused the surge in graphics memory and caused OOM, but this issue was resolved by downgrading 4.34.1

It's OK to not know what caused it (this is what the issues are for! :) ) but as mentioned before we have very many requests and need you to help us too if we are to be able to address them at a good pace. 

Could you provide a minimal i.e. as few lines of code as possible, snippet which we can run to replicate the issue? 

hi bro, do you have training problems? I am training llama, with 4.35 loss will be significantly higher than 4.33!

> hi bro, do you have training problems? I am training llama, with 4.35 loss will be significantly higher than 4.33!

yes I do

Wondering that if it is the change of the logic in line 520~522

from

``` python
query_states = query_states.to(torch.float16)
key_states = key_states.to(torch.float16)
value_states = value_states.to(torch.float16)
```

to

``` python
query_states = query_states.to(target_dtype)
key_states = key_states.to(target_dtype)
value_states = value_states.to(target_dtype)
```

Another possible logical estimate that may be relevant is the change in the "attention_mask" section, starting from 4.35, to: "_prepare_4d_causal_attention_mask"

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.

I have encountered the exact same issue with transformers + flash-attn. Any version >= 4.35 would get OOM while the exact same code with transformers < 4.35 won't.

@jedyang97 Thanks for confirming the. Could you provide a minimal reproducer which can replicate this issue?