Does lora caused memory leak in transformers ?
Author: Randolph-zeng
Pull request: no

### System Info

Issues persisted across several peft version 0.3, 0.4, 0.5(dev)
Accelerator I used are 0.21.0,
Pytorch version I tried are 1.13.0 and 2.0 and both experience same memory explosion
transformers == 4.31.0

### Who can help?

@ArthurZucker @younesbelkada @pacman100 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

I am using the llama2 model with the alpaca project and what I found is that by simply switching the load_in_8bit flag on and off(I promise the other variables are kept untouched as the original script), the RAM on cuda jumps from 13GB to 51GBs. https://github.com/tloen/alpaca-lora/blob/main/finetune.py#L114
I think even loading the full model without any half precision should only takes around 30GBs and the tunable params are ~20 millions, the increase of memory usage does not makes sense to me .
To reproduce, please check out the repo and pip install the requirement file as I did today, and run the finetune script with the default datasets and llama2 model. 
The memory usage between the 8bits flag on and off looks like below:
![image](https://github.com/huggingface/transformers/assets/11933185/a2207318-4275-4aea-bc1d-62077db30f58)
![image](https://github.com/huggingface/transformers/assets/11933185/7f2ec560-8c39-4952-b8a9-b1cff2185706)




### Expected behavior

I think the expected behavior should use much less memory tho. Any insight will be appreciated!

Comments:
Hi @Randolph-zeng 
Thanks for the issue, to clarify, do you get 13GB with load_in_8bit or without it?

hi @younesbelkada , the 13GB is with load_in_8bit on , and 51GB off. Is this an expected behavior or I miscalculated the math :(

@younesbelkada  To add more to the picture, I actually took a long route to debug this. At first I thought it is the lora did not loaded correctly and the whole model was loaded and marked as trainable. However, the lora did take effect and the trainable params are indeed around 20 millions. When I use a debugger and torch.profiler, I found that there seems to be a memory leak and the memory consumption increases abnormally everytime it passed the attention layer (20millions of total trainable params across ALL layers should incurs additional 150+MBs of memory usage for EVERY layer). I thought this was a lora issue at the begininig but later I found out the load_in_8bit causes a big difference in the memory usage and thus posted the questions here instead. The lora side confirms to me that the lora module does not have any memory leak ...(https://github.com/huggingface/peft/issues/832#issuecomment-1682356876)

@younesbelkada I just found the same problem of memory leak using lora + llama2-7b

@younesbelkada @pacman100 @ArthurZucker The memory leak seems to be a serious issue. Can we get some insights on what is happening under the hood or maybe some hotfix we can apply by ourselves in the meantime ? Any help will be greatly appreciated !  

@Randolph-zeng Thanks for your interest! I am building a multi-modal LLM with LLaMA-2 and LoRA; thus, my setting may differ slightly. Yesterday, I found that each time my code ran past the following forward function, the GPU memory increased by ~200MB, finally leading to CUDA OOM.
```
outputs = self.llama_model(
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            return_dict=True,
            labels=target_ids,
        )
```
After debugging, I tried to detach and put all resulting tensors to the CPU and then clean the cuda empty cache, which solved this problem. Note that all these discussed above are during my model inference instead of training. (Yes, my model also need to call forward() instead of generate() during inference)

However, the `load_in_8bit` seems still has errors. the forward() will return nan if `load_in_8bit` is `True`. So I just turn this option off temporarily.

Hi @zhuole1025 ! Thanks for your response and solution. And yes, I did observe similar increase pattern of memory usage when calling the forward, which is why I suspect there is some memory leak. 
Load in 8 bit does solve the memory leak issue in my case which I think might due to the different settings of loading models and keeping reference of related variables. I think by manually detaching and clearing cache as you do should also resolve the issue. However, in my training case and my current lib versions, simply load in 8bit already does the work :)
I think this is actually a serious issue that worth the attention of the HF team, hope they can have some spare time to look at this .

We have also experienced memory leaks in TRL that we have managed to solve by calling 

```python
import gc

gc.collect()
torch.cuda.empty_cache()
gc.collect()

```
After each training step: https://github.com/huggingface/trl/blob/main/trl/core.py#L266-L268

But not sure if it is related to PEFT + LoRA

Hello @Randolph-zeng,

There is no memory leakage. Let's take things step by step.

When using PEFT with INT8/INT4 Quantization, by default gradient checkpointing is enabled. 

When using PEFT alone, let's put some numbers on the expected GPU memory usage:
1. Model in FP32 would take 4bytes per params*7B params = 28GB.
2. Let's add the space by LoRA params = 4bytes per params*0.022B params = 0.088GB
3. Activations take up a huge GPU memory for larger batches and longer sequences. Refer paper [Reducing Activation Recomputation in Large Transformer Models](https://arxiv.org/pdf/2205.05198.pdf) for approximate derivation of memory occupied by activations. Below is the screenshot of the relevant snippet showing activation memory per layer:
![Screenshot 2023-08-22 at 1 32 08 PM](https://github.com/huggingface/transformers/assets/13534540/3e311f89-e3af-4538-8439-480d4668eb65). 
The above formula considers half-precision, so we need to approx double it for FP32. Applying it to your code leveraging micro_batch_size=4, hidden_dimension= 4096, sequence_length=256, num_layers=32, num_attention_heads=32 and adjusting the above formula to consider the activations of LoRA layers which takes `2(4sb(r+h)+sbh)` where lora_rank=8: it will be:
```
(2*(256*4*4096(34+((5*32*256)/4096)))+2*(256*4*4(8+4096)+256*4*4096)) * 32 layers = 12.25GB
```
4. Optimizer state and gradients: 
Gradients in FP32 will take: 0.088GB same as the lora params
Optimizer state: 0.176GB
5. Total approx GPU memory usage: 28+0.088+12.25+0.088+0.176 = 40.602 GB
6. Code: https://github.com/pacman100/DHS-LLM-Workshop/tree/main/personal_copilot/training. Running on single A100 80Gb GPU with the below command:
```
python train.py  \
 --model_path "meta-llama/Llama-2-7b-chat-hf"   \
--dataset_name "smangrul/hf-stack-v1"   \
--subset "data"   \
--data_column "content"   \
--split "train"   \
--seq_length 256   \
--max_steps 2000  \
 --batch_size 4   \
--gradient_accumulation_steps 1   \
--learning_rate 5e-5   \
--lr_scheduler_type "cosine"   \
--weight_decay 0.01   \
--num_warmup_steps 30   \
--eval_freq 100   \
--save_freq 500   \
--log_freq 25   \
--num_workers 4   \
--no_fp16   \
--output_dir "delete-me-llama7b-personal-copilot-A100-80GB"   \
--fim_rate 0.5  \
 --fim_spm_rate 0.5   \
--use_peft_lora   \
--lora_r 8  \
 --lora_alpha 32   \
--lora_dropout 0.1   \
--lora_target_modules "q_proj,v_proj" \
--no_gradient_checkpointing
```

output logs:
```
Size of the train set: 5875. Size of the validation set: 30
100%|████████████████████████████████████████████████████████████████████████████████████████| 400/400 [00:10<00:00, 37.52it/s]
The character to token ratio of the dataset is: 2.88
FIM is not supported by tokenizer, disabling FIM
FIM is not supported by tokenizer, disabling FIM
Loading the model
Downloading shards: 100%|████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00,  9.91it/s]
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████| 2/2 [00:03<00:00,  1.65s/it]
Downloading (…)neration_config.json: 100%|████████████████████████████████████████████████████| 188/188 [00:00<00:00, 1.40MB/s]
PeftModelForCausalLM(
  (base_model): LoraModel(
    (model): LlamaForCausalLM(
      (model): LlamaModel(
        (embed_tokens): Embedding(32000, 4096)
        (layers): ModuleList(
          (0-31): 32 x LlamaDecoderLayer(
            (self_attn): LlamaAttention(
              (q_proj): Linear(
                in_features=4096, out_features=4096, bias=False
                (lora_dropout): ModuleDict(
                  (default): Dropout(p=0.1, inplace=False)
                )
                (lora_A): ModuleDict(
                  (default): Linear(in_features=4096, out_features=8, bias=False)
                )
                (lora_B): ModuleDict(
                  (default): Linear(in_features=8, out_features=4096, bias=False)
                )
                (lora_embedding_A): ParameterDict()
                (lora_embedding_B): ParameterDict()
              )
              (k_proj): Linear(in_features=4096, out_features=4096, bias=False)
              (v_proj): Linear(
                in_features=4096, out_features=4096, bias=False
                (lora_dropout): ModuleDict(
                  (default): Dropout(p=0.1, inplace=False)
                )
                (lora_A): ModuleDict(
                  (default): Linear(in_features=4096, out_features=8, bias=False)
                )
                (lora_B): ModuleDict(
                  (default): Linear(in_features=8, out_features=4096, bias=False)
                )
                (lora_embedding_A): ParameterDict()
                (lora_embedding_B): ParameterDict()
              )
              (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
              (rotary_emb): LlamaRotaryEmbedding()
            )
            (mlp): LlamaMLP(
              (gate_proj): Linear(in_features=4096, out_features=11008, bias=False)
              (up_proj): Linear(in_features=4096, out_features=11008, bias=False)
              (down_proj): Linear(in_features=11008, out_features=4096, bias=False)
              (act_fn): SiLUActivation()
            )
            (input_layernorm): LlamaRMSNorm()
            (post_attention_layernorm): LlamaRMSNorm()
          )
        )
        (norm): LlamaRMSNorm()
      )
      (lm_head): Linear(in_features=4096, out_features=32000, bias=False)
    )
  )
)
Starting main loop
Using the `WANDB_DISABLED` environment variable is deprecated and will be removed in v5. Use the --report_to flag to control the integrations used for logging result (for instance --report_to none).
Cloning https://huggingface.co/smangrul/delete-me-llama7b-personal-copilot-A100-80GB into local empty directory.
Training...
{'loss': 1.5655, 'learning_rate': 4.166666666666667e-05, 'epoch': 0.01}                                                        
{'loss': 1.4344, 'learning_rate': 4.998728546500082e-05, 'epoch': 0.03}                                                        
{'loss': 1.3379, 'learning_rate': 4.993565483078743e-05, 'epoch': 0.04}                                                        
{'loss': 1.2521, 'learning_rate': 4.984439542929117e-05, 'epoch': 0.05}                                                        
{'eval_loss': 1.3657349348068237, 'eval_runtime': 101.306, 'eval_samples_per_second': 4.6, 'eval_steps_per_second': 1.155, 'epoch': 0.05}
{'loss': 1.27, 'learning_rate': 4.971365229370284e-05, 'epoch': 0.06}                                                          
{'loss': 1.2335, 'learning_rate': 4.9543633206385834e-05, 'epoch': 0.07}                                                       
  9%|███████▌                                                                               | 174/2000 [06:47<53:44,  1.77s/it]
```

Memory usage:
![Screenshot 2023-08-22 at 2 04 08 PM](https://github.com/huggingface/transformers/assets/13534540/8725624d-9c57-410f-8a85-27a151b20d24)

So, it is around **40GB** as per our calculations above. 

Now, enable gradient checkpoining to decrease the activations VRAM with below command:
```
python train.py  \
 --model_path "meta-llama/Llama-2-7b-chat-hf"   \
--dataset_name "smangrul/hf-stack-v1"   \
--subset "data"   \
--data_column "content"   \
--split "train"   \
--seq_length 256   \
--max_steps 2000  \
 --batch_size 4   \
--gradient_accumulation_steps 1   \
--learning_rate 5e-5   \
--lr_scheduler_type "cosine"   \
--weight_decay 0.01   \
--num_warmup_steps 30   \
--eval_freq 100   \
--save_freq 500   \
--log_freq 25   \
--num_workers 4   \
--no_fp16   \
--output_dir "delete-me-llama7b-personal-copilot-A100-80GB"   \
--fim_rate 0.5  \
 --fim_spm_rate 0.5   \
--use_peft_lora   \
--lora_r 8  \
 --lora_alpha 32   \
--lora_dropout 0.1   \
--lora_target_modules "q_proj,v_proj" 
```

Memory usage:
![Screenshot 2023-08-22 at 2 13 00 PM](https://github.com/huggingface/transformers/assets/13534540/fcd4e7f6-be06-4ab3-b715-703b2ae8e89d)

We can see that GPU memory usage reduced by ~**12GB** (40-28)

But as per the output logs, the training time increased **1.5X** due to recomputation during backward passes:
```
Size of the train set: 5875. Size of the validation set: 30
100%|████████████████████████████████████████████████████████████████████████████████████████| 400/400 [00:10<00:00, 37.54it/s]
The character to token ratio of the dataset is: 2.88
FIM is not supported by tokenizer, disabling FIM
FIM is not supported by tokenizer, disabling FIM
Loading the model
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████| 2/2 [00:01<00:00,  1.25it/s]
PeftModelForCausalLM(
  (base_model): LoraModel(
    (model): LlamaForCausalLM(
      (model): LlamaModel(
        (embed_tokens): Embedding(32000, 4096)
        (layers): ModuleList(
          (0-31): 32 x LlamaDecoderLayer(
            (self_attn): LlamaAttention(
              (q_proj): Linear(
                in_features=4096, out_features=4096, bias=False
                (lora_dropout): ModuleDict(
                  (default): Dropout(p=0.1, inplace=False)
                )
                (lora_A): ModuleDict(
                  (default): Linear(in_features=4096, out_features=8, bias=False)
                )
                (lora_B): ModuleDict(
                  (default): Linear(in_features=8, out_features=4096, bias=False)
                )
                (lora_embedding_A): ParameterDict()
                (lora_embedding_B): ParameterDict()
              )
              (k_proj): Linear(in_features=4096, out_features=4096, bias=False)
              (v_proj): Linear(
                in_features=4096, out_features=4096, bias=False
                (lora_dropout): ModuleDict(
                  (default): Dropout(p=0.1, inplace=False)
                )
                (lora_A): ModuleDict(
                  (default): Linear(in_features=4096, out_features=8, bias=False)
                )
                (lora_B): ModuleDict(
                  (default): Linear(in_features=8, out_features=4096, bias=False)
                )
                (lora_embedding_A): ParameterDict()
                (lora_embedding_B): ParameterDict()
              )
              (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
              (rotary_emb): LlamaRotaryEmbedding()
            )
            (mlp): LlamaMLP(
              (gate_proj): Linear(in_features=4096, out_features=11008, bias=False)
              (up_proj): Linear(in_features=4096, out_features=11008, bias=False)
              (down_proj): Linear(in_features=11008, out_features=4096, bias=False)
              (act_fn): SiLUActivation()
            )
            (input_layernorm): LlamaRMSNorm()
            (post_attention_layernorm): LlamaRMSNorm()
          )
        )
        (norm): LlamaRMSNorm()
      )
      (lm_head): Linear(in_features=4096, out_features=32000, bias=False)
    )
  )
)
Starting main loop
Using the `WANDB_DISABLED` environment variable is deprecated and will be removed in v5. Use the --report_to flag to control the integrations used for logging result (for instance --report_to none).
/home/sourab/DHS-LLM-Workshop/personal_copilot/training/delete-me-llama7b-personal-copilot-A100-80GB is already a clone of https://huggingface.co/smangrul/delete-me-llama7b-personal-copilot-A100-80GB. Make sure you pull the latest changes with `repo.git_pull()`.
Training...
  0%|                                                                                                 | 0/2000 [00:00<?, ?it/s]/home/sourab/miniconda3/envs/hf/lib/python3.10/site-packages/torch/utils/checkpoint.py:426: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.
  warnings.warn(
{'loss': 1.5655, 'learning_rate': 4.166666666666667e-05, 'epoch': 0.01}                                                        
{'loss': 1.4344, 'learning_rate': 4.998728546500082e-05, 'epoch': 0.03}                                                        
  3%|██▏                                                                                   | 51/2000 [02:14<1:25:06,  2.62s/it]
```


For further decrease in VRAM and increasing speed, use Flash Attention V2. FOr further decrease in VRAM of model, use half-precision/8-bit quantization/4-bit quantization. 

Hope this helps.


Hi @pacman100 , thanks a lot for taking the time to reproduce and provide such detailed explaination on memory consumption. While I fully agree on your examples and analysis, I am afraid there are some subtle differences on our setup that might make us miss the real bug here:
My footprint of 51GB of memory usage is not caused by a full precision loading of a 7B model, it is the result of switching load_in_8bits in the line https://github.com/tloen/alpaca-lora/blob/main/finetune.py#L114. This makes the transformers load in FP16 precision and takes around 14~15 GB as expected. However, when the training starts, the memory consumption jumps to 51 GB. All I have done is to switch the flag load_in_8bits and the memory usage changes so dramatically that can not be explained by the formula that works in normal cases.

I think the memory leak did happen and can be partially supported by the gc and clear cache operations @younesbelkada  and @zhuole1025 suggested above. Since you already have the llama model downloaded, if I am not asking too much from you, can you please checkout the lora-alpaca repo and run the scripst with load_in_8bit on and off ? I am worried that we missed a really important bug here. 


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 been having similar issues.

Please refer to https://github.com/dvlab-research/LongLoRA/blob/be3bb473b4ca2c291c1d26419f6e409d774dd422/supervised-fine-tune.py#L314-L316.

I experienced exactly the same issue, and resolve it by just adding one line `model.gradient_checkpointing_enable()` on 8bit turned off.