a memory leak in evaluation
Author: nooblyh
Pull request: no

## Environment info

- `transformers` version: 4.4.0.dev0
- Platform: Linux-3.10.0-957.el7.x86_64-x86_64-with-glibc2.10
- Python version: 3.8.8
- PyTorch version (GPU?): 1.7.1+cu101 (True)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: True
- Using distributed or parallel set-up in script?: False


### Who can help

@sgugger 

## Information

Model I am using (Bert, XLNet ...): 
albert-base-v2 but use a hidden_size of 2048 and a num_attention_heads of 16, distilled from albert-xlarge-v2.

The problem arises when using:
* [x] the official example scripts: (give details below)
examples/text-classification/run_glue.py

The tasks I am working on is:
* [x] an official GLUE/SQUaD task: (give the name)
GLUE QQP task

## To reproduce

Steps to reproduce the behavior:

I want to evaluate my model on the GLUE QQP task. If I don't use eval_accumulation_step, my GPUs are OOM. But if I use eval_accumulation_step and my memory usage will grow up to the memory limit (>250GB) until the first process is killed due to this issue. So I assumed that maybe there is a memory leak.
My running script is as below.
```
CUDA_VISIBLE_DEVICES=0 ~/.conda/envs/thesis-lyh/bin/python run_glue.py \
  --model_name_or_path $MODEL_PATH \
  --task_name $TASK_NAME \
  --eval_accumulation_step 1 \
  --do_eval \
  --max_seq_length 128 \
  --per_device_eval_batch_size 1 \
  --output_dir output/glue/$TASK_NAME/$MODEL_NAME/
```
No matter what batch_size  and accumulation_step are set to, the above problem still occurs. 
But I am doing fine in models hosted in the model hub and a smaller model I distilled in the same way.

## Expected behavior
I have 250GB RAM so it should be enough to save the result.


Comments:
Which model are you using? There is no reason the predictions for QQP should OOM even your GPU, unless the model is outputting more than the logits.

Thank you very much for your reply! My model config is as below:
```JSON
{
  "architectures": [
    "AlbertForMaskedLM"
  ],
  "attention_probs_dropout_prob": 0,
  "bos_token_id": 2,
  "classifier_dropout_prob": 0.1,
  "down_scale_factor": 1,
  "embedding_size": 128,
  "eos_token_id": 3,
  "gap_size": 0,
  "hidden_act": "gelu_new",
  "hidden_dropout_prob": 0,
  "hidden_size": 2048,
  "initializer_range": 0.02,
  "inner_group_num": 1,
  "intermediate_size": 3072,
  "layer_norm_eps": 1e-12,
  "max_position_embeddings": 512,
  "model_type": "albert",
  "net_structure_type": 0,
  "num_attention_heads": 16,
  "num_hidden_groups": 1,
  "num_hidden_layers": 12,
  "num_memory_blocks": 0,
  "pad_token_id": 0,
  "type_vocab_size": 2,
  "vocab_size": 30000
}
```

And I'm loading my model like this:
```Python
state_dict = torch.load(os.path.join(model_args.model_name_or_path, "checkpoint.pth"))
model = AutoModelForSequenceClassification.from_pretrained(
    pretrained_model_name_or_path=None,
    config=config,
    state_dict=state_dict,
    use_auth_token=True if model_args.use_auth_token else None
)
```

And these are all the named parameters:
```
albert.embeddings.position_ids   torch.Size([1, 512])
albert.embeddings.word_embeddings.weight         torch.Size([30000, 128])
albert.embeddings.position_embeddings.weight     torch.Size([512, 128])
albert.embeddings.token_type_embeddings.weight   torch.Size([2, 128])
albert.embeddings.LayerNorm.weight       torch.Size([128])
albert.embeddings.LayerNorm.bias         torch.Size([128])
albert.encoder.embedding_hidden_mapping_in.weight        torch.Size([2048, 128])
albert.encoder.embedding_hidden_mapping_in.bias          torch.Size([2048])
albert.encoder.albert_layer_groups.0.albert_layers.0.full_layer_layer_norm.weight        torch.Size([2048])
albert.encoder.albert_layer_groups.0.albert_layers.0.full_layer_layer_norm.bias          torch.Size([2048])
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.query.weight      torch.Size([2048, 2048])
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.query.bias        torch.Size([2048])
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.key.weight        torch.Size([2048, 2048])
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.key.bias          torch.Size([2048])
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.value.weight      torch.Size([2048, 2048])
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.value.bias        torch.Size([2048])
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.dense.weight      torch.Size([2048, 2048])
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.dense.bias        torch.Size([2048])
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.LayerNorm.weight          torch.Size([2048])
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.LayerNorm.bias    torch.Size([2048])
albert.encoder.albert_layer_groups.0.albert_layers.0.ffn.weight          torch.Size([3072, 2048])
albert.encoder.albert_layer_groups.0.albert_layers.0.ffn.bias    torch.Size([3072])
albert.encoder.albert_layer_groups.0.albert_layers.0.ffn_output.weight   torch.Size([2048, 3072])
albert.embeddings.LayerNorm.weight       torch.Size([128])                                                                                                                        
albert.embeddings.LayerNorm.bias         torch.Size([128])                                                                                                                                  
albert.encoder.embedding_hidden_mapping_in.weight        torch.Size([2048, 128])                                                                                                            
albert.encoder.embedding_hidden_mapping_in.bias          torch.Size([2048])                                                                                                                 
albert.encoder.albert_layer_groups.0.albert_layers.0.full_layer_layer_norm.weight        torch.Size([2048])                                                                                 
albert.encoder.albert_layer_groups.0.albert_layers.0.full_layer_layer_norm.bias          torch.Size([2048])                                                                                 
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.query.weight      torch.Size([2048, 2048])                                                                                   
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.query.bias        torch.Size([2048])                                                                                         
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.key.weight        torch.Size([2048, 2048])                                                                                   
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.key.bias          torch.Size([2048])                                                                                         
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.value.weight      torch.Size([2048, 2048])                                                                                   
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.value.bias        torch.Size([2048])                                                                                         
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.dense.weight      torch.Size([2048, 2048])                                                                                   
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.dense.bias        torch.Size([2048])                                                                                         
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.LayerNorm.weight          torch.Size([2048])                                                                                 
albert.encoder.albert_layer_groups.0.albert_layers.0.attention.LayerNorm.bias    torch.Size([2048])                                                                                         
albert.encoder.albert_layer_groups.0.albert_layers.0.ffn.weight          torch.Size([3072, 2048])                                                                                           
albert.encoder.albert_layer_groups.0.albert_layers.0.ffn.bias    torch.Size([3072])                                                                                                         
albert.encoder.albert_layer_groups.0.albert_layers.0.ffn_output.weight   torch.Size([2048, 3072])                                                                                           
albert.encoder.albert_layer_groups.0.albert_layers.0.ffn_output.bias     torch.Size([2048])                                                                                                 
albert.pooler.weight     torch.Size([2048, 2048])                                                                                                                                           
albert.pooler.bias       torch.Size([2048])                                                                                                                                                 
classifier.weight        torch.Size([2, 2048])                                                                                                                                              
classifier.bias          torch.Size([2])  
```
The training and evaluation code is run_glue.py.

By the way, my model is pretrained(distilled) in a distributed manner(distributedDataParallel). I'm wondering if it is ok to run GLUE tasks this way? I will be grateful for any help you can provide. @sgugger 

I find these logs and I guess this is why my training and evaluation failed.
```
[INFO|trainer.py:472] 2021-04-04 23:14:56,386 >> The following columns in the training set  don't have a corresponding argument in `AlbertForSequenceClassification.forward` and have been ignored: question1, question2, idx.
[INFO|trainer.py:472] 2021-04-04 23:14:56,389 >> The following columns in the evaluation set  don't have a corresponding argument in `AlbertForSequenceClassification.forward` and have been ignored: question1, question2, idx.
```
But I am confused about what happened inside the trainer.

The fact the model has been trained in a distributed manner is not relevant and shouldn't impact this. The warning you get is also not related and normal if you're running the `run_glue` script: it's just informing you that the `Trainer` is dropping those columns after the preprocessing, since they is no model argument matching.

I'm trying to reproduce but everything is working fine on my side. If you just use a randomly initialized ALBERT with this config do you have the same problem? (I can run evaluation without problem on my side for that)

Thank you for your reply. I double-check my config today and find that I am reusing the config from distillation and the output_hidden_states is set to true......I am very sorry for my carelessness and thank you so much for your time and attention.

Ah I understand better now :-)