Potential Memory Leakage during inference using DistilBert/Bert
Author: TOP-RX
Pull request: no

### System Info

transformer: 4.24.0
python: 3.8.13

### 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, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Hello, 
I am potentially facing the memory leakage problem when using DistilBert or Bert to do inference, the following is my code:
```
   @torch.no_grad()
    def inference(self, hidden_state, mask, id):
        self.eval()
        print("begin", torch.cuda.max_memory_allocated("cuda:4")/ 1024 / 1024)
        distilbert_output = self.bert(inputs_embeds=hidden_state, attention_mask=mask, return_dict=False)
        print("middle", torch.cuda.max_memory_allocated("cuda:4")/ 1024 / 1024)
        hidden_state = distilbert_output[0]                    
        pooled_output = hidden_state[:, 0] 
        x = pooled_output
        x = F.dropout(x, p=args.dropout, training=self.training)

        del hidden_state, pooled_output
        
        print("final", torch.cuda.max_memory_allocated("cuda:4")/ 1024 / 1024)
        for i, lin in enumerate(self.lins[:-1]):
            x = lin(x)
            x = F.relu(x)
            x = F.dropout(x, p=args.dropout, training=self.training)
        x = self.lins[-1](x)
        print("final2", torch.cuda.max_memory_allocated("cuda:4")/ 1024 / 1024)
        torch.cuda.empty_cache()
        return x
```
And the result of each printing of memory usage is:
begin : 4200
middle : 38000
final : 38000
final 2 : 38000

It seems `distilbert_output = self.bert(inputs_embeds=hidden_state, attention_mask=mask, return_dict=False)` will not release the memory as MLP does, there is a huge increase between "begin" and "middle", but no increase between "final" and "final 2". Could I have some ideas about this issue? Thanks.

### Expected behavior

GPU should release memory after transformer inference 

Comments:
cc @ArthurZucker and @younesbelkada 

Hi @TOP-RX 
Can you try to add the `torch.cuda_empty_cache()` before `"final"`? Also you might need to combine it with 
```python
import gc

gc.collect()
```
After that call. See this comment from a previous issue for reference: https://github.com/huggingface/transformers/issues/21094#issuecomment-1396951333

Hello @sgugger @younesbelkada @ArthurZucker,

Thanks for your reply! I tried to include the `torch.cuda.empty_cache()` before "final" and also include `gc.collect()` whether before or after `torch.cuda.empty_cache()`, the issue still happened. And I also use `print(torch.cuda.memory_allocated("cuda:4")/ 1024 / 1024)` to check the allocated memory:

```
    @torch.no_grad()
    def inference(self, hidden_state, mask, id):
        self.eval()
        print("begin_max", torch.cuda.max_memory_allocated("cuda:4")/ 1024 / 1024)
        print("begin", torch.cuda.memory_allocated("cuda:4")/ 1024 / 1024)
        distilbert_output = self.bert(inputs_embeds=hidden_state, attention_mask=mask, return_dict=False)
        print("middle_max", torch.cuda.max_memory_allocated("cuda:4")/ 1024 / 1024)
        print("middle", torch.cuda.memory_allocated("cuda:4")/ 1024 / 1024)

        gc.collect()
        
        torch.cuda.empty_cache()

        gc.collect()

        hidden_state = distilbert_output[0]                    
        pooled_output = hidden_state[:, 0] 
        x = pooled_output
        x = F.dropout(x, p=args.dropout, training=self.training)

        del hidden_state, pooled_output
        
        print("final_max", torch.cuda.max_memory_allocated("cuda:4")/ 1024 / 1024)
        print("final", torch.cuda.memory_allocated("cuda:4")/ 1024 / 1024)
        for i, lin in enumerate(self.lins[:-1]):
            x = lin(x)
            #x = self.bns[i](x)
            x = F.relu(x)
            x = F.dropout(x, p=args.dropout, training=self.training)
        x = self.lins[-1](x)
        self.z_mlp = self.z_mlp.to(device)
        self.z_mlp[id] = x.clone().detach()
        print("final2", torch.cuda.max_memory_allocated("cuda:4")/ 1024 / 1024)
        torch.cuda.empty_cache()
        return x
```
Here are the results:

begin_max 4217.28662109375
begin 4217.28662109375
middle_max 39844.28662109375
middle 7967.28662109375
final_max 39844.28662109375
final 7996.58349609375

there is also a big gap between the  `max_memory_allocated` and `memory_allocated`, could I have some further advices? Thanks.

I have tired several different ways which I found, but the problem is still existing, is this normal?

So one thing to note, you are using `gc.collect();torch.cuda.empty_cache()` but you the `del hidden_state, pooled_output` is after. You should first delete, then call   `gc.collect();torch.cuda.empty_cache()`. 

Note: regarding ` max_memory_allocated`

> By default, this returns the peak allocated memory since the beginning of this program. [reset_peak_memory_stats()](https://pytorch.org/docs/stable/generated/torch.cuda.reset_peak_memory_stats.html#torch.cuda.reset_peak_memory_stats) can be used to reset the starting point in tracking this metric.

I don't see a problem with those larger values for `max_memory_allocated`. They are just the peak values.

Hello @ydshieh @ArthurZucker,

Thanks for your help! Before I call the function as shown above, I use:

```
num_layers = 1  # Specify the number of layers to remove
encoder_layers = model.bert.transformer.layer[-num_layers:]
model.bert.transformer.layer = nn.ModuleList(encoder_layers)
```

to control the number of transformer layers in my model. However, if I include `@torch.no_grad()` in the function I showed above, here are the results:

> 1 transformer layer: allocated memory: 800MB, max allocated: 2300MB
> 2 transformer layers: allocated memory: 840MB, max allocated: 2560MB

if I just comment out the `@torch.no_grad()` in the above code to do a comparison:

> 1 transformer layer: allocated memory: 4107MB, max allocated: 4299MB
> 2 transformer layers: allocated memory: 7564MB, max allocated: 7756MB

For the first case with `@torch.no_grad()`, we don't need to store the intermediate value for backward, it's reasonable the GPU memory is less than the second case. In the second case, the GPU usage is proportional to the number of layers I used which is consistent with my intuition. What makes me confused is no matter how many transformer layers I used in the first case(with `@torch.no_grad()`), the GPU memory usage is almost same. I am wondering if I misunderstand something?

Any help would be appreciated!

Hi @TOP-RX 

It's necessary to provide a self-complete code snippet. With only the definition of `def inference` but not the inputs you passed to it, nothing we can help. Also please import everything necessary in the code snippet.

> What makes me confused is no matter how many transformer layers I used in the first case(with @torch.no_grad()),

Could you let us know where do you put the print statements in the function `inference` to get:

> transformer layer: allocated memory
> transformer layers: allocated memory

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.