CUDA OOM in the middle of training when the training data is large
## Environment info
- `transformers` version: 4.6.0.dev0
- Platform: Linux-3.10.0-1160.24.1.el7.x86_64-x86_64-with-centos-7.9.2009-Core
- Python version: 3.7.9
- PyTorch version (GPU?): 1.8.0 (False)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help
 @sgugger @patrickvonplaten, @patil-suraj

## Information
I am using https://github.com/huggingface/transformers/blob/9e147d31f67a03ea4f5b11a5c7c3b7f8d252bfb7/examples/seq2seq/run_seq2seq.py to train MT5/base on custom parallel data. The code works well when the training data is <=100K but throws CUDA out of memory error in the middle of training when I train with 200K (or beyond) data. 


The error message is here:
loading configuration file https://huggingface.co/google/mt5-base/resolve/main/config.json from cache at /scratch/st-amuham01-1/ganeshjw/projects/mt5_beluga/cache_dir/5ebfd830555547194403d6803baa127970de59b443c04b7a1a60b16a97ed3958.3950cd4aaa701cb6f55a976ff996001a5fb09bbbe7ba9084619949d9016f519e
Model config MT5Config {
  "_name_or_path": "/home/patrick/hugging_face/t5/mt5-base",
  "architectures": [q
    "T5ForConditionalGeneration"
  ],
  "d_ff": 2048,
  "d_kv": 64,
  "d_model": 768,
  "decoder_start_token_id": 0,
  "dropout_rate": 0.1,
  "eos_token_id": 1,
  "feed_forward_proj": "gated-gelu",
  "initializer_factor": 1.0,
  "is_encoder_decoder": true,
  "layer_norm_epsilon": 1e-06,
  "model_type": "mt5",
  "num_decoder_layers": 12,
  "num_heads": 12,
  "num_layers": 12,
  "output_past": true,
  "pad_token_id": 0,
  "relative_attention_num_buckets": 32,
  "tie_word_embeddings": false,
  "tokenizer_class": "T5Tokenizer",
  "transformers_version": "4.6.0.dev0",
  "use_cache": true,
  "vocab_size": 250112
}

loading configuration file https://huggingface.co/google/mt5-base/resolve/main/config.json from cache at /scratch/st-amuham01-1/ganeshjw/projects/mt5_beluga/cache_dir/5ebfd830555547194403d6803baa127970de59b443c04b7a1a60b16a97ed3958.3950cd4aaa701cb6f55a976ff996001a5fb09bbbe7ba9084619949d9016f519e
Model config MT5Config {
  "_name_or_path": "/home/patrick/hugging_face/t5/mt5-base",
  "architectures": [
    "T5ForConditionalGeneration"
  ],
  "d_ff": 2048,
  "d_kv": 64,
  "d_model": 768,
  "decoder_start_token_id": 0,
  "dropout_rate": 0.1,
  "eos_token_id": 1,
  "feed_forward_proj": "gated-gelu",
  "initializer_factor": 1.0,
  "is_encoder_decoder": true,
  "layer_norm_epsilon": 1e-06,
  "model_type": "mt5",
  "num_decoder_layers": 12,
  "num_heads": 12,
  "num_layers": 12,
  "output_past": true,
  "pad_token_id": 0,
  "relative_attention_num_buckets": 32,
  "tie_word_embeddings": false,
  "tokenizer_class": "T5Tokenizer",
  "transformers_version": "4.6.0.dev0",
  "use_cache": true,
  "vocab_size": 250112
}

Can't load following files from cache: ['added_tokens_file', 'tokenizer_file'] and cannot check if these files are necessary for the tokenizer to operate.
loading file https://huggingface.co/google/mt5-base/resolve/main/spiece.model from cache at /scratch/st-amuham01-1/ganeshjw/projects/mt5_beluga/cache_dir/4764ec347af4d2d6286acbe1d9d630ac0afd8554a4c4a64170e0b663fd2e2412.84ea7af2df68dc8db434d3160aab65cce8ac63ce5b6f7743f8c9a4a14b4f77e2
loading file https://huggingface.co/google/mt5-base/resolve/main/special_tokens_map.json from cache at /scratch/st-amuham01-1/ganeshjw/projects/mt5_beluga/cache_dir/0d7d5b3fc19bf58d4b274990c8bcf5e307726bc18d95f40a1436dfb6a0892f85.294ebaa4cd17bb284635004c92d2c4d522ec488c828dcce0c2471b6f28e3fe82
loading file https://huggingface.co/google/mt5-base/resolve/main/tokenizer_config.json from cache at /scratch/st-amuham01-1/ganeshjw/projects/mt5_beluga/cache_dir/afba33be693521ccefbde6d03b93b5c517d7108ba31f6c08000ed52c2cea45c9.28bbf90ae7962b1b7211c0ce8b2006f968c82439ec9c47e0847ba63642f9435a
loading weights file https://huggingface.co/google/mt5-base/resolve/main/pytorch_model.bin from cache at /scratch/st-amuham01-1/ganeshjw/projects/mt5_beluga/cache_dir/3b7e8056d4ed71d8d7ac2dea78627c4be77ed136399c05b563d4116abfcd9418.1afec9001b62cd5a347e7fd4b664e503ca2377606e11b9ddb8ec1d7b79bc3952
All model checkpoint weights were used when initializing MT5ForConditionalGeneration.

All the weights of MT5ForConditionalGeneration were initialized from the model checkpoint at google/mt5-base.
If your task is similar to the task the model of the checkpoint was trained on, you can already use MT5ForConditionalGeneration for predictions without further training.
100%|██████████| 196/196 [00:46<00:00,  4.22ba/s]
100%|██████████| 1/1 [00:00<00:00, 11.27ba/s]
100%|██████████| 1/1 [00:00<00:00,  6.31ba/s]
***** Running training *****
  Num examples = 195996
  Num Epochs = 5
  Instantaneous batch size per device = 8
  Total train batch size (w. parallel, distributed & accumulation) = 32
  Gradient Accumulation steps = 1
  Total optimization steps = 30625
  0%|          | 0/30625 [00:00<?, ?it/s]/arc/project/st-amuham01-1/ganeshjw/.conda/envs/blink/lib/python3.7/site-packages/torch/nn/parallel/_functions.py:65: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.
  warnings.warn('Was asked to gather along dimension 0, but all '
 19%|█▉        | 5906/30625 [55:18<4:09:46,  1.65it/s]Traceback (most recent call last):
  File "/scratch/st-amuham01-1/ganeshjw/projects/mt5_beluga/run_seq2seq_general.py", line 879, in <module>
    main()
  File "/scratch/st-amuham01-1/ganeshjw/projects/mt5_beluga/run_seq2seq_general.py", line 625, in main
    train_result = trainer.train(resume_from_checkpoint=None)
  File "/arc/project/st-amuham01-1/ganeshjw/.conda/envs/blink/lib/python3.7/site-packages/transformers-4.6.0.dev0-py3.7.egg/transformers/trainer.py", line 1192, in train
    tr_loss += self.training_step(model, inputs)
  File "/arc/project/st-amuham01-1/ganeshjw/.conda/envs/blink/lib/python3.7/site-packages/transformers-4.6.0.dev0-py3.7.egg/transformers/trainer.py", line 1590, in training_step
    loss = self.compute_loss(model, inputs)
  File "/arc/project/st-amuham01-1/ganeshjw/.conda/envs/blink/lib/python3.7/site-packages/transformers-4.6.0.dev0-py3.7.egg/transformers/trainer.py", line 1622, in compute_loss
    outputs = model(**inputs)
  File "/arc/project/st-amuham01-1/ganeshjw/.conda/envs/blink/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/arc/project/st-amuham01-1/ganeshjw/.conda/envs/blink/lib/python3.7/site-packages/torch/nn/parallel/data_parallel.py", line 167, in forward
    outputs = self.parallel_apply(replicas, inputs, kwargs)
  File "/arc/project/st-amuham01-1/ganeshjw/.conda/envs/blink/lib/python3.7/site-packages/torch/nn/parallel/data_parallel.py", line 177, in parallel_apply
    return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
  File "/arc/project/st-amuham01-1/ganeshjw/.conda/envs/blink/lib/python3.7/site-packages/torch/nn/parallel/parallel_apply.py", line 86, in parallel_apply
    output.reraise()
  File "/arc/project/st-amuham01-1/ganeshjw/.conda/envs/blink/lib/python3.7/site-packages/torch/_utils.py", line 429, in reraise
    raise self.exc_type(msg)
RuntimeError: Caught RuntimeError in replica 0 on device 0.
Original Traceback (most recent call last):
  File "/arc/project/st-amuham01-1/ganeshjw/.conda/envs/blink/lib/python3.7/site-packages/torch/nn/parallel/parallel_apply.py", line 61, in _worker
    output = module(*input, **kwargs)
  File "/arc/project/st-amuham01-1/ganeshjw/.conda/envs/blink/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/arc/project/st-amuham01-1/ganeshjw/.conda/envs/blink/lib/python3.7/site-packages/transformers-4.6.0.dev0-py3.7.egg/transformers/models/t5/modeling_t5.py", line 1505, in forward
    return_dict=return_dict,
  File "/arc/project/st-amuham01-1/ganeshjw/.conda/envs/blink/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/arc/project/st-amuham01-1/ganeshjw/.conda/envs/blink/lib/python3.7/site-packages/transformers-4.6.0.dev0-py3.7.egg/transformers/models/t5/modeling_t5.py", line 959, in forward
    output_attentions=output_attentions,
  File "/arc/project/st-amuham01-1/ganeshjw/.conda/envs/blink/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/arc/project/st-amuham01-1/ganeshjw/.conda/envs/blink/lib/python3.7/site-packages/transformers-4.6.0.dev0-py3.7.egg/transformers/models/t5/modeling_t5.py", line 638, in forward
    output_attentions=output_attentions,
  File "/arc/project/st-amuham01-1/ganeshjw/.conda/envs/blink/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/arc/project/st-amuham01-1/ganeshjw/.conda/envs/blink/lib/python3.7/site-packages/transformers-4.6.0.dev0-py3.7.egg/transformers/models/t5/modeling_t5.py", line 545, in forward
    output_attentions=output_attentions,
  File "/arc/project/st-amuham01-1/ganeshjw/.conda/envs/blink/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/arc/project/st-amuham01-1/ganeshjw/.conda/envs/blink/lib/python3.7/site-packages/transformers-4.6.0.dev0-py3.7.egg/transformers/models/t5/modeling_t5.py", line 502, in forward
    attn_weights, p=self.dropout, training=self.training
  File "/arc/project/st-amuham01-1/ganeshjw/.conda/envs/blink/lib/python3.7/site-packages/torch/nn/functional.py", line 1076, in dropout
    return _VF.dropout_(input, p, training) if inplace else _VF.dropout(input, p, training)
RuntimeError: CUDA out of memory. Tried to allocate 96.00 MiB (GPU 0; 31.75 GiB total capacity; 29.60 GiB already allocated; 2.00 MiB free; 30.11 GiB reserved in total by PyTorch)

 19%|█▉        | 5906/30625 [55:20<3:51:35,  1.78it/s]

Any help would be highly appreciated. Thanks.

Hi
I also have observed the same issue with t5-base model  and mt5-small model <|||||>Could you maybe make use of the `group_by_length` training argument? This will put the largest batch first to make sure OOM are detected in the very beginning (best feature ever by @sgugger )<|||||>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/master/CONTRIBUTING.md) are likely to be ignored.