Memory leak when using CLIPTextModel
Author: minsuk00
Pull request: no
Labels: name: Performance

name: Multimodal

### System Info

- `transformers` version: 4.26.1
- Platform: Linux-5.15.0-105-generic-x86_64-with-glibc2.35
- Python version: 3.11.5
- Huggingface_hub version: 0.17.3
- PyTorch version (GPU?): 2.1.2 (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?

- @ArthurZucker 
- @younesbelkada 

### 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 can't free GPU memory after I use CLIPTextModel
Also, memory is allocated in another device for some reason

problem should be reproduced by using the following code snippet
```python
from transformers import CLIPTextModel
import torch

clip_text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14").to("cuda:1")
del clip_text_model
torch.cuda.empty_cache()
```

### Expected behavior

- CLIPTextModel is placed in "cuda:1", but for some reason, memory gets allocated in "cuda:0" when I call torch.cuda.empty_cache()
- Also, memory is still not freed for "cuda:1"
- As a result, both "cuda:0" and "cuda:1" have some memory allocated

I've also tried using garbage collection and explicitly moving model to cpu, but they don't work.

Comments:
Hi @minsuk00 
you can also try the `release_memory` utility method from `accelerate.utils` - cc @muellerzr 

@younesbelkada  -cc @muellerz 
Thanks for the suggestion, but it doesn't seem to work.
`
clip_text_model = accelerate.utils.release_memory(clip_text_model)
`
does not free any GPU memory.

Additionally, calling `clip_text_model.cpu()` or `torch.cuda.empty_cache()` simply results in the behavior described above.

cc @muellerzr regarding the accelerate behaviour. 

Regarding `torch.cuda.empty_cache()` it's recommended that this function is not manually used c.f. [a related issue](https://github.com/huggingface/transformers/pull/31530), and [this discussion in the pytorch forum](https://discuss.pytorch.org/t/about-torch-cuda-empty-cache/34232/2)

Hey there @minsuk00, Hope you are doing well. (Redirected from #33345 to here). I primarily observe two concerns raised regarding the VRAM usages in here. Let's clear these one by one.

## **_Small VRAM Usage Spike on 2nd GPU/Non-primary GPUs_**
When loading a model, PyTorch initializes a CUDA context for managing resources like device memory and kernel execution. This setup includes detecting which GPUs are available. By default, CUDA might allocate a small amount of memory on all available GPUs, not just the one we explicitly specified for model to be deployed on. This can cause a minor memory spike on non-primary GPUs.

```python
clip_text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14").to("cuda:1")
```
The code above places the model on GPU-1. However, without specific settings, CUDA initializes contexts on all GPUs, leading to small memory footprints on each.

To ensure that only GPU-1 is utilized, and to prevent CUDA from initializing on other GPUs, you can set the `CUDA_VISIBLE_DEVICES` environment variable at the beginning of your script, as follows:
```python
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"  # Only GPU-1 is recognized
```
This command restricts CUDA to see and use only the GPU specified by us and  avoides unnecessary memory allocation on others.

🤗 Tip: Always define `CUDA_VISIBLE_DEVICES` before importing `torch`, to ensure that CUDA configurations are applied correctly:
```python
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

import torch

# rest of the code
```

## **_Residual VRAM Usage After Freeing Up Memory_**

When we use PyTorch with CUDA to load a model, the framework initializes a CUDA context. The CUDA could be thought of as a data structure that manages crucial GPU resources such as device memory allocation, kernel execution, and synchronization mechanisms.

So previously, to make an attempt to clear the VRAM entirely, we executed the following code: 
```python
del clip_text_model
gc.collect()
torch.cuda.empty_cache()
```
This frees up most of the VRAM used by the model, but running this does not dismantle the CUDA context or other minor memory allocations made during the model's initialization. As a result, we still observe some residual VRAM usage. The CUDA context, as a persistent entity, maintains the state and resources required for efficient GPU operations across the lifecycle of the program, ensuring the overhead of initializing these resources is not repeatedly incurred.

The good news is that residual VRAM usage is generally consistent across different models, regardless of their type or size. 

For example:
(1)For `CLIPTextModel`: Once the model is loaded with the following config, `CLIPTextModel` consumes `667MB` of VRAM on the second GPU at peak, and at the end of this script after deleting the model, with garbage collection and emptying the cache, we are left with a VRAM usage of `103MB`.
```python
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

import gc
import torch
from transformers import CLIPTextModel

clip_text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14").to("cuda:1")
del clip_text_model
gc.collect()
torch.cuda.empty_cache()
```

And similarly:
(2)For `Meta-Llama-3.1-8B-Instruct`: Once the model is loaded with the following config, `Meta-Llama-3.1-8B-Instruct` consumes `12.2GB` of VRAM on the second GPU at peak, and at the end of this script after deleting the model, with garbage collection and emptying the cache, we are left with a VRAM usage of `103MB`.

```python
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

import gc
import torch
from transformers import AutoModelForCausalLM
checkpoint = "meta-llama/Meta-Llama-3.1-8B-Instruct"
model = AutoModelForCausalLM.from_pretrained(checkpoint, 
                                             device_map="auto",
                                             torch_dtype=torch.float16,
                                             token=hugging_face_token)
del model
gc.collect()
torch.cuda.empty_cache()
```

🌟 **Important point to note**:
If you ran both the above scripts sequentially, the residual VRAM usage would not sum up to `206MB`; it would still be just `103MB`. This indicates that the CUDA context, once initialized, does not duplicate its foundational resource allocations when additional models are loaded in the same session.

Therefore, loading multiple models does not cumulatively increase residual VRAM usage beyond what is necessary for the CUDA context and other minimal allocations. After using one model, you can delete it and load another as needed without concern for having extra VRAM overhead thought to be carried by another model load.

If you are looking to mitigate even this small residual VRAM usage, I have written the next section discussing techniques to do the same. But if you want to continue with loading another model and continue with your work after loading the first model, i have briefly highlighted the guidelines in the [TLDR & Ending Notes](#tldr) section at the end.

## Mitigating Residual VRAM Usage

### **_(1) Restart the IPy-Kernel:_**
Restarting the Jupyter Notebook kernel clears all defined variables and the CUDA context, effectively freeing up all GPU memory used during the session. This is a straightforward method to reset the memory state of the GPU after the model operations:

### **_(2) Using Terminal to Run the Script:_**
Running your scripts directly through a terminal session instead of a notebook can help manage memory more effectively. Each time the script completes, the Python process terminates, which clears all memory allocations including the CUDA context:

Here’s how to set up and run your script:
```python
%%writefile model_load_script.py
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

import torch
from transformers import CLIPTextModel

clip_text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14").to("cuda:1")

# Operations are performed here
# Once completed, clean up the model and VRAM
del clip_text_model
gc.collect()
torch.cuda.empty_cache()
```
To execute the script, run the following command in your terminal:
```bash
python model_load_script.py
```
### **_(3) Using `numba` to Reset the GPU:_**
For situations where restarting the notebook or terminal is not feasible, you can use `numba`'s GPU management features to forcefully reset the GPU. This method clears all memory allocations and should be used with caution as it affects all CUDA kernels and memory on the GPU:

```python
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

import torch
from numba import cuda
from transformers import CLIPTextModel

clip_text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14").to("cuda:1")

# After other operations, clean up the model and VRAM
del clip_text_model
gc.collect()
torch.cuda.empty_cache()
device = cuda.get_current_device()
device.reset()
```
**Note:** Use this method at the end of your workflow, as it will disrupt all ongoing GPU processes and reset the state entirely, similar to rebooting the GPU.

<a id="tldr"></a>
## 🤗 **_TLDR & Ending Notes_**
1. For the first issue, setting `CUDA_VISIBLE_DEVICES` before importing `torch` solves the error of non-primary GPU VRAM spike.

2. For the second issue:
   - **(I)** If you just want to remove residual VRAM usage at the end of the script run, you can:
     - (i) Restart the Python Kernel
     - (ii) Use Terminal to run the script
     - (iii) Use `numba` at the end of the script

   - **(II)** If you want to continue doing some `abc()` function or load another `model_y` and do not need the previously loaded `model_x`:
```python
del model_x
gc.collect()
torch.cuda.empty_cache()
```
The code above will clear the VRAM occupied by `model_x` and can be used by `model_y` or can be used to perform other operations which would require VRAM for their execution.

Hope this clarifies the miscellaneous VRAM spikes and Residual VRAM usage🤗, Let me know if something is unclear.

Hey @minsuk00, does the answer above fixes your issue ? 

@SunMarc @nnilayy 
Hey, apologies for the late response.

### 1.
For the first problem with the error of non-primary GPU VRAM spike, setting the `CUDA_VISIBLE_DEVICES ` to 1 works well, but only when I set the device to "cuda:0" instead of "cuda:1"
```python
clip_text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14").to("cuda:0")
```
Data is correctly allocated at GPU 1 as a result.
Otherwise, when I use ```.to("cuda:1")```, it causes this error:

> RuntimeError: CUDA error: invalid device ordinal
> CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
> For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
> Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

### 2.
I confirmed that indeed, the residual VRAM after freeing the memory is shared among multiple models.

---
Thanks for the help! Really appreciate it.

> For the first problem with the error of non-primary GPU VRAM spike, setting the CUDA_VISIBLE_DEVICES  to 1 works well, but only when I set the device to "cuda:0" instead of "cuda:1"
> 
> clip_text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14").to("cuda:0")
> Data is correctly allocated at GPU 1 as a result.
> Otherwise, when I use .to("cuda:1"), it causes this error:

Yes, you should do `.to("cuda:0")` because setting CUDA_VISIBLE_DEVICES = 1 means that you only see the gpu n*1, so "cuda:0" target this gpu and since there are no other visible gpus, "cuda:1" is not a valid device ordinal. 

If your issue is fixed, feel free to close the PR. Thanks !