Significantly increased VRAM usage for Mixtral qlora training compared to 4.36.2?
Author: DocShotgun
Pull request: no

### System Info

The environment is a Runpod container with python 3.10, single A100 80gb, transformers 4.37.0dev (3cefac1d974db5e2825a0cb2b842883a628be7a0), using axolotl training script (https://github.com/OpenAccess-AI-Collective/axolotl).

### Who can help?

_No response_

### Information

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

Hello, just tried doing a training run on the dev version of transformers (as of 3cefac1d974db5e2825a0cb2b842883a628be7a0) via the common training repository axolotl (https://github.com/OpenAccess-AI-Collective/axolotl) and noticed that I went OOM using the same configuration that I had previously used successfully with transformers 4.36.2 stable. And not even just a small difference - I had to reduce my batch size by 4x to make the training fit in VRAM.

I was previously able to fit 8192 ctx, batch size 4, grad accum steps 2 without difficulty, but I found that I now had to reduce my batch size to 1 to avoid OOM. The relevant training hyperparameters are:
```
load_in_4bit: true
sequence_len: 8192
sample_packing: true
pad_to_sequence_len: true
lora_r: 32
lora_alpha: 16
lora_dropout: 0.05
lora_target_linear: true
optimizer: adamw_bnb_8bit
bf16: true
fp16: false
tf32: true
gradient_checkpointing: true
flash_attention: true

no deepspeed or fsdp
no evals
```
Would appreciate any insights into what caused the massive increase in memory usage. I noticed that ehartford's latest dolphin 2.7 qlora used a batch size of 3 per device at 16k ctx on A100 80gb, so surely I'm missing something here?

### Expected behavior

The training run should take a relatively similar amount of VRAM as it did previously with the same config.

Comments:
Hey! Thanks for the report, here are potential PRs that I would suspect: 
- #28142 which fixes some of the FA2 bugs
- #28061 wich should fix gradient checkpointing issues
Pinging @younesbelkada for when he comes back, would be great if you can isolate the commit that lead to this in the meantime 🤗 

This may not be relevant to you but I found [this recent change](https://github.com/OpenAccess-AI-Collective/axolotl/commit/4d2e842e46bf8bd6dd0fda4d2667a7e7d80b4cd4) to Axolotl has made a significant difference to VRAM usage. Previously I could just squeeze in a LoRA on a 34B model on my 3x3090s at batch size 2, seq length 4096, now it OOMs immediately. I undid the change and it fits again. 

> This may not be relevant to you but I found [this recent change](https://github.com/OpenAccess-AI-Collective/axolotl/commit/4d2e842e46bf8bd6dd0fda4d2667a7e7d80b4cd4) to Axolotl has made a significant difference to VRAM usage. Previously I could just squeeze in a LoRA on a 34B model on my 3x3090s at batch size 2, seq length 4096, now it OOMs immediately. I undid the change and it fits again.

Hmm it's certainly possible since that commit was in between when I did my initial train and the run where I had to drop the batch size. Unfortunately don't have a training instance up right now, so I'd have to test it the next time I try to train.

I've determined that the cause of the increased VRAM usage was indeed axolotl changing the default for use_reentrant to False for gradient checkpointing. Going to go ahead and close the issue.

thanks for sharing the solution! 🤗