CUDA out of memory when use DistillBert for inference and use hidden_state as input_embeds
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, since the memory cost of tuning the whole distillbert is still large for me, I am trying to get the last hidden states of a modified distillbert which only has first 5 layers, then use it as the input for a  distillbert which only has the last 1 transformer layers and fine tune it (only fine tune the last transformer layer to save memory). **But when I try to get the hidden_states, it will always be out of memory whatever batch size I use.** Since I can use my current GPU to run even Bertseqclassfiication, I think the GPU memory is sufficient. Could I have some help about this?

Besides, **I am not sure if I can use hidden_state directly as input_embeds for another distillbert (in my case, it's a one transformer layer distillbert)** I want to fine tune, here is my current code, please let me know if I am not correct:

```
with open('/X.all.txt', "r") as fin:
    node_text_list = fin.readlines()

model_path = r'/distilbert-base-uncased'
tokenizer = DistilBertTokenizer.from_pretrained(model_path)
X = [tokenizer(text, padding='max_length', max_length=32,
                truncation=True, return_tensors="pt")
        for text in node_text_list]`

input_ids = []
attention_mask = []
for i in range(len(X)):
   input_ids.append(X[i]['input_ids'])
   attention_mask.append(X[i]['attention_mask'])

input_ids = torch.stack(input_ids).squeeze(1)
attention_mask = torch.stack(attention_mask).squeeze(1)

data_set = torch.utils.data.TensorDataset(
            input_ids, attention_mask
         )
# Load the first BERT model
pre_loader = torch.utils.data.DataLoader(
data_set, batch_size=128, num_workers=0, pin_memory=True
)

model_pretrain = DistilBertModel.from_pretrained(args.pretrain)
model_pretrain.transformer.layer = model_pretrain.transformer.layer[:5]

hidden_states_list = []
model_pretrain = model_pretrain.to(device)
for param in model_pretrain.parameters():
     param.requires_grad = False

# Pass the input through the first BERT model
for input_id, mask in pre_loader:
    with torch.no_grad():
        input_id = input_id.to(device)
        mask = mask.to(device)
        outputs = model_pretrain(input_id, mask, return_dict=False)  ###out of memory here
        hidden_states = outputs[0]

    # Remove the memory usage of the first BERT model
    del outputs

    # Append the hidden states to the list
    hidden_states_list.append(hidden_states)

# Concatenate the hidden states along the batch dimension
pretrain_hidden_states = torch.cat(hidden_states_list, dim=0)



model = DistilBertModel.from_pretrained(pretrain)

# Remove unnecessary layers from BERT
num_removed_layers = 1  # Specify the number of layers to remove
encoder_layers = model.transformer.layer[-num_removed_layers:]
model.transformer.layer = nn.ModuleList(encoder_layers)

### new sampling
train_set = torch.utils.data.TensorDataset(
            pretrain_hidden_states, data.attention_mask
         )

train_loader = torch.utils.data.DataLoader(
train_set, batch_size=128, num_workers=0, pin_memory=True
)

for hidden_state, mask in train_loader:
     distilbert_output = model(inputs_embeds=hidden_state, attention_mask=mask, return_dict=False)
     hidden_state = distilbert_output[0]                    
     pooled_output = hidden_state[:, 0] 
     x = pooled_output
    ### classification layers
```

### Expected behavior

CUDA should not be out of memory

Comments:
Hi @TOP-RX, 

A few things from first looking at your script: 

* tokenizers already work on batches. There's no need to pass line by line and then concatenate

```python
encoder_inputs = tokenizer(node_text_list, max_length=32, truncation=True, return_tensors="pt")
input_ids = encoder_inputs["input_ids"]
attention_mask = encoder_inputs["attention_mask"]
```

* Creating the dataset like this means that ALL of your data is read into memory (although not GPU) and converted to pytorch arrays at once. Consider loading using `datasets` some [info here](https://huggingface.co/docs/datasets/nlp_load) and tokenizing as a map function applied to the dataset. e.g. [like here](https://huggingface.co/docs/datasets/nlp_process). I highly recommend looking at the [scripts in `examples`](https://github.com/huggingface/transformers/tree/main/examples/pytorch) to see how best to structure training pipelines.  

* Batch size
You mention it fails at any batch size - is this true for batch_size=1? The batch size in this script (128) is quite large

* Increasing memory usage
Does this happen if you just do one forward pass i.e. passing a single batch with no for loop? 
I'd guess some of the memory issues are coming from `hidden_states_list`, which increases in size and the tensors I believe are still on the cuda device. 

Hi @amyeroberts ,

Thanks so much for your reply and advices! I was following your suggestion and found the issue happened with `hidden_states_list.append(hidden_states)` even I used batch size =1, the GPU memory accumulates for each batch to cause the out of memory problem after several batches.

I am also wondering if I could have some suggestions from you about another issue(ignore some unrelated parts):
```
class Net(nn.Module):
     def forward(hidden_state, mask):
     distilbert_output = self.distilbert(inputs_embeds=hidden_state, attention_mask=mask, return_dict=False)
     hidden_state = distilbert_output[0]                    
     pooled_output = hidden_state[:, 0] 
     x = pooled_output
    ### classification layers

# Remove unnecessary layers from BERT
model = Net()
num_removed_layers = 1  # Specify the number of layers to remove
encoder_layers = model.distilbert.transformer.layer[-num_removed_layers:]
model.distilbert.transformer.layer = nn.ModuleList(encoder_layers)
```
1. Since I just want to use previous layers as encoder not tuning, and only tune the last layer (as code shown, I remove the layers which were used to generate the hidden states)plus my own model to save memory, Is this a correct way to directly use hidden_state I got as inputs_embeds for a Bert/DistillBert? And it seems if I use input_ids is fine but when I use `inputs_embeds=hidden_state` with the same settings, it will be out of memory. 
2. I found no matter how many layers I removed, the GPU memory usage is almost same, it this normal?

Thanks so much!

@TOP-RX 

> I was following your suggestion and found the issue happened with hidden_states_list.append(hidden_states) even I used batch size =1

OK, this indicates the issue is a result of the script and not code relating to transformers. 
Questions about debugging custom training objectives or scripts are best placed in our [forums](https://discuss.huggingface.co/). We try to reserve the github issues for feature requests and bug reports.

Regarding your questions: 

1. Same as above - this is a question for our forums :) 
2. Without more specifics about the model, GPU utilisation vs. layers and what "almost the same" means, it's not possible to help. If you suspect this is a bug, then please open another separate issue giving information about how to reproduce and expected behaviour. 

You are collecting the outputs (here `hidden_states_list`) for a whole dataset (e.g. `for input_id, mask in pre_loader:`). This is not going to work well. You should find a way to save it to some storage like disk with some tools (probably there is some in torch).

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.