t5-large model OOM with FP16, but run well without FP16
Author: Agoniii
Pull request: no

## 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! -->

Copy-and-paste the text below in your GitHub issue and FILL OUT the two last points.

- `transformers` version: 4.17.0
- Platform: Linux-5.4.0-91-generic-x86_64-with-glibc2.10
- Python version: 3.8.12
- PyTorch version (GPU?): 1.11.0a0+17540c5 (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, V100(16GB)
- Using distributed or parallel set-up in script?: Yes, 2 GPUs


### Who can help
@stas00


## Information

Model I am using (Bert, XLNet ...): T5-large

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:
* [x] an official GLUE/SQUaD task: (give the name)
* [ ] my own task or dataset: (give details below)

## To reproduce
I refered to https://github.com/huggingface/transformers/issues/8771#issuecomment-759248400
| Method     | max BS |
| ---      | ---       |
| baseline |  2 (4 will OOM)  |
| fp16 |   1 will OOM |

Steps to reproduce the behavior:
```
git clone https://github.com/huggingface/transformers
cd transformers/examples/legacy/seq2seq
wget https://cdn-datasets.huggingface.co/translation/wmt_en_ro.tar.gz
tar -xzvf wmt_en_ro.tar.gz
```
Baseline script:
```
export BS=2; rm -r output_dir; \
PYTHONPATH=../../src USE_TF=0 python -m torch.distributed.launch \
--nproc_per_node=2 ./finetune_trainer.py --model_name_or_path t5-large \
--output_dir output_dir --adam_eps 1e-06 --data_dir wmt_en_ro --do_eval \
--do_train --evaluation_strategy=steps --freeze_embeds --label_smoothing 0.1 \
--learning_rate 3e-5 --logging_first_step --logging_steps 1000 \
--max_source_length 128 --max_target_length 128 --num_train_epochs 1 \
--overwrite_output_dir --per_device_eval_batch_size \
--per_device_train_batch_size --predict_with_generate --eval_steps 25000 \
--sortish_sampler --task translation_en_to_ro --test_max_target_length 128 \
--val_max_target_length 128 --warmup_steps 500 --n_train 2000 --n_val 500
```
FP16 script:
```
export BS=1; rm -r output_dir; \
PYTHONPATH=../../src USE_TF=0 python -m torch.distributed.launch \
--nproc_per_node=2 ./finetune_trainer.py --model_name_or_path t5-large \
--output_dir output_dir --adam_eps 1e-06 --data_dir wmt_en_ro --do_eval \
--do_train --evaluation_strategy=steps --freeze_embeds --label_smoothing 0.1 \
--learning_rate 3e-5 --logging_first_step --logging_steps 1000 \
--max_source_length 128 --max_target_length 128 --num_train_epochs 1 \
--overwrite_output_dir --per_device_eval_batch_size \
--per_device_train_batch_size --predict_with_generate --eval_steps 25000 \
--sortish_sampler --task translation_en_to_ro --test_max_target_length 128 \
--val_max_target_length 128 --warmup_steps 500 --n_train 2000 --n_val 500 \
--fp16
```
<!-- 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.-->
```
Traceback (most recent call last):
  File "./finetune_trainer.py", line 367, in <module>
    main()
  File "./finetune_trainer.py", line 304, in main
    train_result = trainer.train(
  File "/opt/conda/lib/python3.8/site-packages/transformers/trainer.py", line 1400, in train
    tr_loss_step = self.training_step(model, inputs)
  File "/opt/conda/lib/python3.8/site-packages/transformers/trainer.py", line 1984, in training_step
    loss = self.compute_loss(model, inputs)
  File "/workspace/huggingface/transformers/examples/legacy/seq2seq/seq2seq_trainer.py", line 180, in compute_loss
    loss, _ = self._compute_loss(model, inputs, labels)
  File "/workspace/huggingface/transformers/examples/legacy/seq2seq/seq2seq_trainer.py", line 173, in _compute_loss
    logits = model(**inputs, use_cache=False)[0]
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 930, in forward
    output = self.module(*inputs[0], **kwargs[0])
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/transformers/models/t5/modeling_t5.py", line 1635, in forward
    decoder_outputs = self.decoder(
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/transformers/models/t5/modeling_t5.py", line 1030, in forward
    layer_outputs = layer_module(
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/transformers/models/t5/modeling_t5.py", line 717, in forward
    hidden_states = self.layer[-1](hidden_states)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/transformers/models/t5/modeling_t5.py", line 327, in forward
    forwarded_states = self.DenseReluDense(forwarded_states)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/transformers/models/t5/modeling_t5.py", line 285, in forward
    hidden_states = self.wi(hidden_states)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/linear.py", line 103, in forward
    return F.linear(input, self.weight, self.bias)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/functional.py", line 1971, in linear
    return torch._C._nn.linear(input, weight, bias)
RuntimeError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 1; 15.78 GiB total capacity; 14.58 GiB already allocated; 7.75 MiB free; 14.69 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
  9%|▉         | 89/1000 [00:31<05:20,  2.84it/s]
ERROR:torch.distributed.elastic.multiprocessing.api:failed (exitcode: 1) local_rank: 0 (pid: 2260) of binary: /opt/conda/bin/python
```

(edited by @stas00 to wrap long cmd lines to make it easier to read the Issue)

Comments:
1. Before we can proceed, any chance you could port your code to the modern version of finetune_trainer - this code is very old and is no longer being tested or supported and is available only for those who still have to use it.

   I remember there were a few stages to the transition, I tried to keep track of the changes here: https://github.com/huggingface/transformers/issues/10036 if it helps.

2. In general fp16 mixed precision doesn't always save memory over straight fp32, since you actually need to have even more memory to allocate both fp16 and fp32 weights, and then some operations use less memory under amp than fp32. You will find the full explanation and breakdown here: https://huggingface.co/docs/transformers/performance#fp16

   If you want to use a bigger batch size use deepspeed with cpu memory offload. You can see the stats in the comment you linked to.

Thanks @stas00, I will try to the new code.

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.