GPU memory usage increased continuously in validation
Author: sdlee130
Pull request: no

### System Info

- `transformers` version: 4.30.2
- Platform: Linux-5.15.0-78-generic-x86_64-with-glibc2.29
- Python version: 3.8.10
- Huggingface_hub version: 0.15.1
- Safetensors version: 0.3.1
- PyTorch version (GPU?): 2.0.1+cu117 (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?

@sgugger   
I was trying to finetune a model that inherits from BertForSequenceClassification.

```python
model = CustomForSequenceClassification.from_pretrained(model_name,
                                                        output_hidden_states=True,
                                                        return_dict=False,
                                                        num_labels=2,
                                                       )
```   
`CustomForSeuenceClassification` is a class that inherit from BertForSequenceClassification.
I used `return_dict=False` and `output_hidden_states=True` to get all encoders' output.
But, if I use `output_hidden_states=True`, GPU memory usage increased during validation process. 
This issue is a problem when training on a large dataset. 
When I set `output_hidden_states` to False, that issue didn't occured. 
I am using GLUE dataset to train and validation.
Below code is a class that inherits from BertForSequenceClassfication.
```python
class CustomForSequenceClassification(BertForSequenceClassification):
    
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.config = config

        self.bert = BertModel(config)
        classifier_dropout = (
            config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
        )
        self.dropout = nn.Dropout(classifier_dropout)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)

        self.post_init()
        
    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        token_type_ids: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        last_hidden_states = outputs[1]
        all_hidden_states = outputs[2]
        pooled_output = last_hidden_states

        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)


        loss = None
        if labels is not None:
            if self.config.problem_type is None:
                if self.num_labels == 1:
                    self.config.problem_type = "regression"
                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
                    self.config.problem_type = "single_label_classification"
                else:
                    self.config.problem_type = "multi_label_classification"

            if self.config.problem_type == "regression":
                loss_fct = MSELoss()
                if self.num_labels == 1:
                    loss = loss_fct(logits.squeeze(), labels.squeeze())
                else:
                    loss = loss_fct(logits, labels)
            elif self.config.problem_type == "single_label_classification":
                loss_fct = CrossEntropyLoss()
                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
            elif self.config.problem_type == "multi_label_classification":
                loss_fct = BCEWithLogitsLoss()
                loss = loss_fct(logits, labels)
        if not return_dict:
            output = (logits,) + outputs[2:]
            return ((loss,) + output) if loss is not None else output

        return SequenceClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )
```   

### Information

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

### Tasks

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

### Reproduction

```python
 #1. import      
import math
from transformers import BertTokenizer, BertForSequenceClassification, AdamW

import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from transformers import TrainingArguments, Trainer, TrainerCallback
from datasets import load_dataset
from torch import nn

import pdb
import ipdb
from typing import List, Optional, Tuple, Union

from transformers.modeling_outputs import SequenceClassifierOutput
from transformers import BertModel, BertLayer

import math
import os
import warnings
import random
from dataclasses import dataclass
from typing import List, Optional, Tuple, Union

import torch
import torch.utils.checkpoint
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from sklearn.metrics import accuracy_score, precision_recall_fscore_support, matthews_corrcoef
from scipy.stats import pearsonr, spearmanr

# 2. load model, tokenizer, dataset, compute_metrics
model_name = 'bert-base-uncased'
tokenizer = BertTokenizer.from_pretrained(model_name)

train = load_dataset('glue', 'qqp', split='train')
valid = load_dataset('glue', 'qqp', split='validation')

class BERT_Dataset(torch.utils.data.Dataset):
    def __init__(self, dataset1, dataset2, label, tokenizer): 
        self.dataset1 = dataset1
        self.dataset2 = dataset2
        self.label = label
        self.tokenizer = tokenizer

    def __getitem__(self, idx):
        text1 = self.dataset1[idx]
        text2 = self.dataset2[idx]
        tokens = self.tokenizer(text1, text2,
                               max_length=128,
                               padding="max_length",
                               truncation=True,
                               )
        tokens['label'] = torch.LongTensor([self.label[idx]])
        
        return tokens

    def __len__(self):
        return len(self.label)

train_dataset = BERT_Dataset(train['question1'],
                             train['question2'],
                             train['label'],
                             tokenizer,
                            ) 
valid_dataset = BERT_Dataset(valid['question1'],
                             valid['question2'],
                             valid['label'],
                             tokenizer,
                            )

def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    predictions = np.argmax(predictions[0], axis=1)
    _, _, f1, _ = precision_recall_fscore_support(labels, predictions, average='binary')
    acc = accuracy_score(labels, predictions)
    return {
        'f1': f1,
        'accuracy': acc,
    }

# 3. inherit from BertForSequenceClassification
class CustomForSequenceClassification(BertForSequenceClassification):
    
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.config = config

        self.bert = BertModel(config)
        #self.attention_layer = BertLayer(config)
        #self.pooler = BertPooler(config)
        #self.attn = Attention(hidden_dim=config.hidden_size, method='dot')
        classifier_dropout = (
            config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
        )
        self.dropout = nn.Dropout(classifier_dropout)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)

        self.post_init()
        
    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        token_type_ids: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        last_hidden_states = outputs[1]
        all_hidden_states = outputs[2]
        pooled_output = last_hidden_states

        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)


        loss = None
        if labels is not None:
            if self.config.problem_type is None:
                if self.num_labels == 1:
                    self.config.problem_type = "regression"
                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
                    self.config.problem_type = "single_label_classification"
                else:
                    self.config.problem_type = "multi_label_classification"

            if self.config.problem_type == "regression":
                loss_fct = MSELoss()
                if self.num_labels == 1:
                    loss = loss_fct(logits.squeeze(), labels.squeeze())
                else:
                    loss = loss_fct(logits, labels)
            elif self.config.problem_type == "single_label_classification":
                loss_fct = CrossEntropyLoss()
                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
            elif self.config.problem_type == "multi_label_classification":
                loss_fct = BCEWithLogitsLoss()
                loss = loss_fct(logits, labels)
        if not return_dict:
            output = (logits,) + outputs[2:]
            return ((loss,) + output) if loss is not None else output

        return SequenceClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

# 4. generate model object

model = CustomForSequenceClassification.from_pretrained(model_name,
                                                        output_hidden_states=False,
                                                        return_dict=False,
                                                        num_labels=2,
                                                       )

# 5. set TrainerArguements

training_ars = TrainingArguments(
    output_dir="./checkpoint",
    num_train_epochs=20,
    per_device_train_batch_size=128,
    per_device_eval_batch_size=16,
    learning_rate=1e-5,
    weight_decay=0.01,
    evaluation_strategy="epoch",
)

trainer = Trainer(
    model=model,
    args=training_ars,
    train_dataset=train_dataset,
    eval_dataset=valid_dataset,
    tokenizer=tokenizer,
    compute_metrics=compute_metrics,
)

# 6. start training

trainer.train()


### Expected behavior

In first validation process, GPU memory usage will increase.
If used dataset is large dataset, then CUDA OOM error will be occurred.

Comments:
Yes, the tensors are accumulated on the GPUs and you are asking for a lot of them with `output_attentions=True`, so it's perfectly normal to see a GPU memory increase. To manage your GPU RAM you will need to set `eval_accumulation_steps` so that the tensors accumulated are offloaded to the CPU every `eval_accumulation_steps` steps/

I already try to use `eval_accumulation_steps`. But when I used that, RAM(not GPU memory) usage increased continuously and then kernel was dead. 

And I have another question. In my knowledge, loss and logit also accumulate in gpu but why GPU memory usage dosen't change? I think it is strange that the GPU memory usage did not change during the validation process when `output_attentions=False`, but it did change when `output_attentions=True`.

Again with `output_attentions=True`, you are asking to get all the attentions outputs of all layers for all inputs. It will take a lot of memory, either on the GPU or on the CPU (if using `eval_accumulation_steps`). Do not ask the model to output those if you do not have the memory to store them.

Thank you for your answering