CUDA out of memory in Seq2SeqTrainer class 
Author: kritika121
Pull request: no
Labels: name: bug

### System Info

```shell
- `transformers` version: 4.11.0
- Platform: Linux-4.14.252-131.483.amzn1.x86_64-x86_64-with-glibc2.9
- Python version: 3.6.13
- PyTorch version (GPU?): 1.10.2+cu102 (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?: No
```


### Who can help?

@sgugger @patrickvonplaten
Hello! I am trying to finetune "sshleifer/distill-pegasus-xsum-16-4" model for a seq2seq2 generation task(specifically summarization) on my own custom dataset(~1800 training data points) using hugging face transformers Seq2SeqTrainer but encountered CUDA OOM error.

I am trying to follow the [finetune-summarization notebook](https://github.com/huggingface/notebooks/blob/main/examples/summarization.ipynb) mentioned by @sgugger 

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Libraries
```bash
import transformers
from datasets import load_dataset, load_metric
from transformers import AutoTokenizer
from transformers import AutoModelForSeq2SeqLM, DataCollatorForSeq2Seq, Seq2SeqTrainingArguments, Seq2SeqTrainer
import nltk
import numpy as np
```

Data
```bash
data_files = {
    "train": "data/train.jsonl",
    "validation": "data/val.jsonl"
}

raw_datasets = load_dataset('json', data_files=data_files)
```

Load tokenizer and model
```bash
model_checkpoint = 'sshleifer/distill-pegasus-xsum-16-4'
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)
```

Process Data
```bash
if model_checkpoint in ["t5-small", "t5-base", "t5-larg", "t5-3b", "t5-11b"]:
    prefix = "summarize: "
else:
    prefix = ""
max_input_length = 1024
max_target_length = 128

def preprocess_function(examples):
    inputs = [prefix + doc for doc in examples["document"]]
    model_inputs = tokenizer(inputs, max_length=max_input_length, truncation=True)

    # Setup the tokenizer for targets
    with tokenizer.as_target_tokenizer():
        labels = tokenizer(examples["summary"], max_length=max_target_length, truncation=True)

    model_inputs["labels"] = labels["input_ids"]
    return model_inputs
tokenized_datasets = raw_datasets.map(preprocess_function, batched=True)
```

Trainer
```bash
metric = load_metric("rouge")
batch_size = 2
model_name = model_checkpoint.split("/")[-1]
args = Seq2SeqTrainingArguments(
    f"{model_name}-finetuned-xsum",
    evaluation_strategy = "epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    weight_decay=0.01,
    save_total_limit=2,
    num_train_epochs=1,
    predict_with_generate=True,
    fp16=True,
)
data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)

def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
    # Replace -100 in the labels as we can't decode them.
    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
    
    # Rouge expects a newline after each sentence
    decoded_preds = ["\n".join(nltk.sent_tokenize(pred.strip())) for pred in decoded_preds]
    decoded_labels = ["\n".join(nltk.sent_tokenize(label.strip())) for label in decoded_labels]
    
    result = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
    # Extract a few results
    result = {key: value.mid.fmeasure * 100 for key, value in result.items()}
    
    # Add mean generated length
    prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in predictions]
    result["gen_len"] = np.mean(prediction_lens)
    
    return {k: round(v, 4) for k, v in result.items()}

trainer = Seq2SeqTrainer(
    model,
    args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    data_collator=data_collator,
    tokenizer=tokenizer,
    compute_metrics=compute_metrics
)

trainer.train()
```

Error
```bash
The following columns in the training set  don't have a corresponding argument in `PegasusForConditionalGeneration.forward` and have been ignored: summary, document.
***** Running training *****
  Num examples = 1599
  Num Epochs = 1
  Instantaneous batch size per device = 2
  Total train batch size (w. parallel, distributed & accumulation) = 2
  Gradient Accumulation steps = 1
  Total optimization steps = 800
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-33-3435b262f1ae> in <module>
----> 1 trainer.train()

~/anaconda3/envs/python3/lib/python3.6/site-packages/transformers/trainer.py in train(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)
   1310                         tr_loss_step = self.training_step(model, inputs)
   1311                 else:
-> 1312                     tr_loss_step = self.training_step(model, inputs)
   1313 
   1314                 if args.logging_nan_inf_filter and (torch.isnan(tr_loss_step) or torch.isinf(tr_loss_step)):

~/anaconda3/envs/python3/lib/python3.6/site-packages/transformers/trainer.py in training_step(self, model, inputs)
   1837         if self.use_amp:
   1838             with autocast():
-> 1839                 loss = self.compute_loss(model, inputs)
   1840         else:
   1841             loss = self.compute_loss(model, inputs)

~/anaconda3/envs/python3/lib/python3.6/site-packages/transformers/trainer.py in compute_loss(self, model, inputs, return_outputs)
   1871         else:
   1872             labels = None
-> 1873         outputs = model(**inputs)
   1874         # Save past state if it exists
   1875         # TODO: this needs to be fixed and made cleaner later.

~/anaconda3/envs/python3/lib/python3.6/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1100         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1101                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1102             return forward_call(*input, **kwargs)
   1103         # Do not call functions when jit is used
   1104         full_backward_hooks, non_full_backward_hooks = [], []

~/anaconda3/envs/python3/lib/python3.6/site-packages/transformers/models/pegasus/modeling_pegasus.py in forward(self, input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, head_mask, decoder_head_mask, cross_attn_head_mask, encoder_outputs, past_key_values, inputs_embeds, decoder_inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict)
   1391             output_attentions=output_attentions,
   1392             output_hidden_states=output_hidden_states,
-> 1393             return_dict=return_dict,
   1394         )
   1395         lm_logits = self.lm_head(outputs[0]) + self.final_logits_bias

~/anaconda3/envs/python3/lib/python3.6/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1100         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1101                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1102             return forward_call(*input, **kwargs)
   1103         # Do not call functions when jit is used
   1104         full_backward_hooks, non_full_backward_hooks = [], []

~/anaconda3/envs/python3/lib/python3.6/site-packages/transformers/models/pegasus/modeling_pegasus.py in forward(self, input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, head_mask, decoder_head_mask, cross_attn_head_mask, encoder_outputs, past_key_values, inputs_embeds, decoder_inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict)
   1226                 output_attentions=output_attentions,
   1227                 output_hidden_states=output_hidden_states,
-> 1228                 return_dict=return_dict,
   1229             )
   1230         # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput when return_dict=True

~/anaconda3/envs/python3/lib/python3.6/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1100         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1101                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1102             return forward_call(*input, **kwargs)
   1103         # Do not call functions when jit is used
   1104         full_backward_hooks, non_full_backward_hooks = [], []

~/anaconda3/envs/python3/lib/python3.6/site-packages/transformers/models/pegasus/modeling_pegasus.py in forward(self, input_ids, attention_mask, head_mask, inputs_embeds, output_attentions, output_hidden_states, return_dict)
    796                         attention_mask,
    797                         layer_head_mask=(head_mask[idx] if head_mask is not None else None),
--> 798                         output_attentions=output_attentions,
    799                     )
    800 

~/anaconda3/envs/python3/lib/python3.6/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1100         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1101                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1102             return forward_call(*input, **kwargs)
   1103         # Do not call functions when jit is used
   1104         full_backward_hooks, non_full_backward_hooks = [], []

~/anaconda3/envs/python3/lib/python3.6/site-packages/transformers/models/pegasus/modeling_pegasus.py in forward(self, hidden_states, attention_mask, layer_head_mask, output_attentions)
    320             attention_mask=attention_mask,
    321             layer_head_mask=layer_head_mask,
--> 322             output_attentions=output_attentions,
    323         )
    324         hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)

~/anaconda3/envs/python3/lib/python3.6/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1100         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1101                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1102             return forward_call(*input, **kwargs)
   1103         # Do not call functions when jit is used
   1104         full_backward_hooks, non_full_backward_hooks = [], []

~/anaconda3/envs/python3/lib/python3.6/site-packages/transformers/models/pegasus/modeling_pegasus.py in forward(self, hidden_states, key_value_states, past_key_value, attention_mask, layer_head_mask, output_attentions)
    207             # self_attention
    208             key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
--> 209             value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
    210 
    211         if self.is_decoder:

~/anaconda3/envs/python3/lib/python3.6/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1100         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1101                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1102             return forward_call(*input, **kwargs)
   1103         # Do not call functions when jit is used
   1104         full_backward_hooks, non_full_backward_hooks = [], []

~/anaconda3/envs/python3/lib/python3.6/site-packages/torch/nn/modules/linear.py in forward(self, input)
    101 
    102     def forward(self, input: Tensor) -> Tensor:
--> 103         return F.linear(input, self.weight, self.bias)
    104 
    105     def extra_repr(self) -> str:

~/anaconda3/envs/python3/lib/python3.6/site-packages/torch/nn/functional.py in linear(input, weight, bias)
   1846     if has_torch_function_variadic(input, weight, bias):
   1847         return handle_torch_function(linear, (input, weight, bias), input, weight, bias=bias)
-> 1848     return torch._C._nn.linear(input, weight, bias)
   1849 
   1850 

RuntimeError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 14.76 GiB total capacity; 13.65 GiB already allocated; 11.75 MiB free; 13.65 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
```

The error can be re-produced on loading any open-source summarization dataset. `
raw_datasets = load_dataset("xsum")
`




### Expected behavior

```shell
Finetune the summarization model.
```


Comments:
Hey @kritika121,

It seems like you don't have enough GPU memory to run your training. Do you maybe have access to a bigger GPU? Otherwise you can try reducing the batch_size, enabling [gradient_checkpointing](https://huggingface.co/docs/transformers/main/en/performance#gradient-checkpointing) or training in [fp16](https://huggingface.co/docs/transformers/main/en/performance#gradient-checkpointing) to save memory.

Thanks @patrickvonplaten I have batch_size of 2 and fp16 is set too. I tried enabling gradient_checkpointing and it worked for me. Thanks for your help!! 

Closing the issue.