Trainer accumulates GPU usage at the beginning of each step
Author: ZL92
Pull request: no

Hello,

My problem is that GPU usage gets increased at the beginning of each step. Although the usage gets decreased with the help of torch.cuda.empty_cache() and gc.collector() during training, OOM errors happened after a while.

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.5.1
- Platform: colab
- Python version: 3.7.10
- PyTorch version (GPU?): 1.8.1+cu101
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No


### Who can help

@sgugger @patrickvonplaten 

<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- albert, bert, xlm: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @Rocketknight1

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

Model hub:

- for issues with a model report at https://discuss.huggingface.co/ and tag the model's creator.

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->

## Information

Models I am using: wav2vec2 and MBart

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

1. Build e2e_model with the following classes: wav2vec2_learn_repr and e2emodel
2. Feed audio and translation under the requirement of the following data_collator.
3. Model training with Trainer.

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

The code for reproducing the error:
```
class wav2vec2_learn_repr(Wav2Vec2PreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.wav2vec2 = Wav2Vec2Model(config)
        self.dropout = nn.Dropout(config.final_dropout)
        self.collapse = collapse_layer
        self.init_weights = ()

    def forward(self,
                input_values,
                attention_mask=None,
                output_attentions=None,
                output_hidden_states=None,
                return_dict=None,
                labels=None):

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        outputs = self.wav2vec2(
            input_values,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        hidden_states = outputs[0]
        hidden_states = self.dropout(hidden_states)
        collapsed_embeddings, attention_masks=cal_collapse_embeddings(hidden_states)
        show_gpu(f'In wav2vec2_learn_repr before del')
        del outputs, hidden_states
        show_gpu(f'In wav2vec2_learn_repr after del')
        torch.cuda.empty_cache()
        show_gpu(f'In wav2vec2_learn_repr empty cache')
        gc.collect()
        show_gpu(f'In wav2vec2_learn_repr gc.collect()')
        return collapsed_embeddings, attention_masks
```
```
class e2emodel(PreTrainedModel):
  def __init__(self,
               wav2vec2_name = "facebook/wav2vec2-large-xlsr-53",  
               mbart_model_name = 'facebook/mbart-large-50-many-to-many-mmt', 
               ):
    super().__init__(PretrainedConfig())
    self.wav2vec2_repr_model = wav2vec2_learn_repr.from_pretrained(wav2vec2_name)
    self.mbart_model = MBartForConditionalGeneration.from_pretrained(mbart_model_name)
    self.wav2vec2_repr_model.to(device)
    self.mbart_model.to(device)
  
  def forward(self,
                input_ids,
                attention_mask=None,
                output_attentions=None,
                output_hidden_states=None,
                return_dict=None,
                labels=None):
    

    show_gpu(f'At start')
    torch.cuda.empty_cache()
    show_gpu(f'empty cache')
    gc.collect()
    show_gpu(f'gc.collect()')
    
    input_ids.to(device)
    # print(f' inputs devices: {input_ids.device}, {labels.device}')
    show_gpu(f'load input_ids')


    collapsed_embeddings, attention_masks = self.wav2vec2_repr_model(input_ids)

    print('collapsed_embeddings, attention_masks', collapsed_embeddings.device, attention_masks.device)
    show_gpu(f'after wav2vec2')
    del input_ids
    show_gpu(f'delete input_ids')
    torch.cuda.empty_cache()
    show_gpu(f'empty cache')
    gc.collect()
    show_gpu(f'gc.collect()')
    labels.to(device)
    show_gpu(f'load mbart inputs')

    output = self.mbart_model(inputs_embeds = collapsed_embeddings, attention_mask = attention_masks, labels = labels)

    show_gpu(f'after mbart')
    del collapsed_embeddings, attention_masks, labels
    show_gpu(f'delete mbart inputs')
    torch.cuda.empty_cache()
    show_gpu(f'empty cache')
    return output
```
```
def data_collator(data):
  translation = [d['translation'] for d in data]
  input_features = [{'input_values': get_inputs_values_from_audio_path(feature_extractor, d['path'])} for d in data]
  #TODO: remove empty audio and its translation
  wav2vec2_inputs = feature_extractor.pad(input_features,
                                padding=True,
                                max_length=None,
                                pad_to_multiple_of=None,
                                return_tensors="pt", )
  batch={}
  batch['inputs_embeds'], batch['attention_mask'] = wav2vec2_learn_repr(wav2vec2_inputs['input_values']) # size [batch_size, nr_sample, 1024]
  with tokenizer.as_target_tokenizer():
    batch['labels'] = tokenizer([d['translation']for d in data], return_tensors='pt', padding=True).input_ids
  return batch
```
```
import torchaudio
resampler = torchaudio.transforms.Resample(orig_freq=48000, new_freq=16000)
def get_inputs_values_from_audio_path(processor, path: str):
    signal, sr = torchaudio.load(main_path + '{}/clips/'.format(src_lang) + path)
    signal = signal.squeeze(0)
    d = (signal.shape[0]/sr)
    resampler.orig_freq = sr
    signal=resampler.forward(signal).numpy()
    input_values = processor(signal, sampling_rate=resampler.new_freq).input_values
    return input_values.tolist()[0]
```
```
import gc
import subprocess
def show_gpu(msg):
    """
    ref: https://discuss.pytorch.org/t/access-gpu-memory-usage-in-pytorch/3192/4
    """
    def query(field):
        return(subprocess.check_output(
            ['nvidia-smi', f'--query-gpu={field}',
                '--format=csv,nounits,noheader'], 
            encoding='utf-8'))
    def to_int(result):
        return int(result.strip().split('\n')[0])
    
    used = to_int(query('memory.used'))
    total = to_int(query('memory.total'))
    pct = used/total
    print('\n' + msg, f'{100*pct:2.1f}% ({used} out of {total})') 
```
## Expected behavior
<!-- A clear and concise description of what you would expect to happen. -->

Here is the GPU usage history from step 7. 

**At start 86.9% (14149 out of 16280)**
empty cache 54.3% (8835 out of 16280)
gc.collect() 54.3% (8835 out of 16280)
load input_ids 54.3% (8835 out of 16280)
In wav2vec2_learn_repr before del 56.7% (9231 out of 16280)
In wav2vec2_learn_repr after del 56.7% (9231 out of 16280)
In wav2vec2_learn_repr empty cache 56.7% (9229 out of 16280)
In wav2vec2_learn_repr gc.collect() 56.7% (9229 out of 16280)
collapsed_embeddings, attention_masks cuda:0 cuda:0
after wav2vec2 56.7% (9229 out of 16280)
delete input_ids 56.7% (9229 out of 16280)
empty cache 56.7% (9229 out of 16280)
gc.collect() 56.7% (9229 out of 16280)
load mbart inputs 56.7% (9229 out of 16280)
after mbart 64.3% (10473 out of 16280)
delete mbart inputs 64.3% (10473 out of 16280)
empty cache 64.3% (10473 out of 16280)
**At start 85.5% (13925 out of 16280)**
empty cache 54.3% (8835 out of 16280)
gc.collect() 54.3% (8835 out of 16280)
load input_ids 54.3% (8835 out of 16280)
In wav2vec2_learn_repr before del 58.9% (9593 out of 16280)
In wav2vec2_learn_repr after del 58.9% (9593 out of 16280)
In wav2vec2_learn_repr empty cache 58.9% (9593 out of 16280)
In wav2vec2_learn_repr gc.collect() 58.9% (9593 out of 16280)
collapsed_embeddings, attention_masks cuda:0 cuda:0
after wav2vec2 58.9% (9593 out of 16280)
delete input_ids 58.9% (9593 out of 16280)
empty cache 58.9% (9593 out of 16280
gc.collect() 58.9% (9593 out of 16280)
load mbart inputs 58.9% (9593 out of 16280)
after mbart 66.7% (10853 out of 16280)
delete mbart inputs 66.7% (10853 out of 16280)
empty cache 66.7% (10853 out of 16280)
**At start 95.4% (15529 out of 16280)**
empty cache 94.6% (15393 out of 16280)
gc.collect() 94.6% (15393 out of 16280)
load input_ids 94.6% (15393 out of 16280)
RuntimeError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 15.90 GiB total capacity; 14.96 GiB already allocated; 21.75 MiB free; 15.00 GiB reserved in total by PyTorch)



Comments:
I've experienced the same issue.