text: 
  
  TITLE
  CUDA out of memory in Seq2SeqTrainer class 
  
  COMMENTS
  0
  
  REACTIONS
  +1: 0
  -1: 0
  laugh: 0
  hooray: 0
  heart: 0
  rocket: 0
  eyes: 0
  
  BODY
  ### 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.
  ```
  