CUDA OOM with increased max input length
Author: karthsu
Pull request: no

### System Info

- `transformers` version: 4.34.0.dev0
- Platform: Linux-5.15.0-76-generic-x86_64-with-glibc2.35
- Python version: 3.10.6
- Huggingface_hub version: 0.16.4
- Safetensors version: 0.3.3
- Accelerate version: not installed
- Accelerate config: not found
- PyTorch version (GPU?): 2.0.1+cu118 (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?: Yes

### Who can help?

_No response_

### 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

I used [this](https://gist.github.com/younesbelkada/9f7f75c94bdc1981c8ca5cc937d4a4da) script to finetune a CodeLlama-34B model on a 4xA100 80GB setup. The PEFT method used here was QLoRA. When i train with a max input length to 4k, the training goes fine. But when I increase the max input length to 8k, the GPU usage is as follows: The first GPU builds up to about ~74GB and then goes OOM shortly after while the other 3 GPUs have a usage of about ~20GB per GPU. I am setting `device_map='auto'` when loading the model in my script.

### Expected behavior

Ideally i expect somewhat equal usage of GPU memory when training and for it not to be blocked on a single GPUs memory. I have no clue why so much memory is being used (when I am using QLoRA especially). Why is this happening and what can be done to remedy this ?

Comments:
cc @younesbelkada 

Hi @karths8 

To increase the sequence length for training, you might be interested to use `BetterTransformer` API from transformers & optimum to use memory efficient attention and increase your maximum sequence length for training.

To enable it, first install optimum package `pip install optimum`
then call `model.to_bettertransformer()` before passing the model to the SFTTrainer object. 
Note that `BetterTransformer` is not compatible when it comes to training with padd tokens, therefore you have to use `packing=True` (i.e. activate this arg: https://gist.github.com/younesbelkada/9f7f75c94bdc1981c8ca5cc937d4a4da#file-finetune_llama_v2-py-L104) in order to concatenate the input sentences all together, and separate them with an EOS token between each sample. That way no padding token is used during training and the model will attend to all tokens during training. 

You can also push that to the next level by force-dispatching the `torch.scaled_dot_product_attention()` to call the Flash attention kernel using the trick presented here: https://twitter.com/younesbelkada/status/1696478075721302143?s=20 / you just need to add the `with torch.backends.cuda.sdp_kernel(enable_math=False, enable_mem_efficient=False, enable_flash=True):` context manager on top of this call: https://gist.github.com/younesbelkada/9f7f75c94bdc1981c8ca5cc937d4a4da#file-finetune_llama_v2-py-L224    

You can more generally read about BT+decoder models here: https://huggingface.co/docs/transformers/perf_infer_gpu_one#decoder-models

In the near future we're also trying to support Flash Attention-2 through a possible native integration https://github.com/huggingface/transformers/pull/25598 , if that PR gets merged, you'll be able to fine-tune your model on any dataset (including non-packed dataset) and with possibly an even larger sequence length

Hope this helps!

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.

Closing as I consider the issue as solved, feel free to re-open if you have more questions

this fix is not compatible with DataCollatorForCompletionOnlyLM  due to "You can use the DataCollatorForCompletionOnlyLM to train your model on the generated prompts only. Note that this works only in the case when packing=False." I have the same problem using 'trl import SFTTrainer'