retain_graph=True required when using a custom BigBird-based model 
Author: helena-balabin
Pull request: no

## Environment info

- `transformers` version: 4.10.2
- Platform: Linux-3.10.0-1160.25.1.el7.x86_64-x86_64-with-glibc2.10
- Python version: 3.8.8
- PyTorch version (GPU?): 1.9.1 (True)
- Tensorflow version (GPU?): 2.5.0 (False)
- 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?: Single device, multi-GPU set up with 4x NVIDIA A100s (40GB) (I'm also using gradient accumulation on top) 

### Who can help

Models:

- Not entirely sure who's familiar with BigBird, ... maybe @patrickvonplaten? 

Library:

- Since it could be related to the Trainer ... also @sgugger?

## Information

Model I am using (Bert, XLNet ...): Custom model based on BigBirdForPreTraining

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)

## Error Message

My custom model class contains the following forward function: 

```
class ProtSTonKGsForPreTraining(BigBirdForPreTraining):
    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        masked_lm_labels=None,
        ent_masked_lm_labels=None,
        prot_masked_lm_labels=None,
        return_dict=None,
        head_mask=None,
    ):
        """Perform one forward pass for a given sequence of text_input_ids + ent_input_ids + prot_input_ids.

        Due to having more than two parts (and a RoBERTa base in the default BigBird model), the NSP objective is
        omitted in this forward function.

        :param input_ids: Concatenation of text + KG (random walk) + protein sequence embeddings
        :param attention_mask: Attention mask of the combined input sequence
        :param token_type_ids: Token type IDs of the combined input sequence
        :param masked_lm_labels: Masked LM labels for only the text part
        :param ent_masked_lm_labels: Masked entity labels for only the KG part
        :param prot_masked_lm_labels: Masked protein labels for only the protein part
        :param return_dict: Whether the output should be returned as a dict or not
        :param head_mask: Used to cancel out certain heads in the Transformer

        :return: Loss, prediction_logits in a BigBirdForPreTrainingOutputWithPooling format
        """
        # 1. Use the LM backbone to get the pre-trained token embeddings
        # batch x number_text_tokens x hidden_size
        # The first element of the returned tuple from the LM backbone forward() pass is the sequence of hidden states
        text_embeddings = torch.cat(
            [
                self.lm_backbone(input_ids[:, i*(self.kg_start_idx//3): (i+1)*(self.kg_start_idx//3)])[0]
                for i in range(3)
            ],
            dim=1,
        )

        # 2. Use the KG backbone to obtain the pre-trained entity embeddings
        # batch x number_kg_tokens x hidden_size
        ent_embeddings = torch.stack(
            [
                # for each numeric index in the random walks sequence: get the embedding vector from the KG backbone
                torch.stack([self.kg_backbone[i.item()] for i in j])
                # for each example in the batch: get the random walks sequence
                for j in input_ids[:, self.kg_start_idx : self.prot_start_idx]
            ],
        )
        # 3. Use the Prot backbone to obtain the pre-trained entity embeddings
        # batch x number_prot_tokens x hidden_size
        prot_embeddings_original_dim = self.prot_backbone(input_ids[:, self.prot_start_idx:])[0]
        prot_embeddings = self.prot_to_lm_hidden_linear(prot_embeddings_original_dim)

        # Concatenate token, KG and prot embeddings obtained from the LM, KG and prot backbones and cast to float
        # batch x seq_len x hidden_size
        inputs_embeds = (
            torch.cat(
                [
                    text_embeddings,
                    ent_embeddings.to(text_embeddings.device),
                    prot_embeddings.to(text_embeddings.device),
                ],
                dim=1,
            )
            .type(torch.FloatTensor)
            .to(self.device)
        )

        # Get the hidden states from the basic STonKGs Transformer layers
        # batch x seq_len x hidden_size
        outputs = self.bert(
            inputs_embeds=inputs_embeds,
            encoder_attention_mask=attention_mask,
            return_dict=True,
        )
        # batch x seq_len x hidden_size
        sequence_output, pooled_output = outputs[:2]

        # Generate the prediction scores (mapping to text and entity vocab sizes + NSP) for the training objectives
        # prediction_scores = Text MLM, entity "MLM" and protein "MLM" scores
        prediction_scores, _ = self.cls(sequence_output, pooled_output)
        # The custom STonKGsELMPredictionHead returns a triple of prediction scores for tokens, entities,
        # and protein sequences, respectively
        (
            token_prediction_scores,
            entity_predictions_scores,
            prot_predictions_scores,
        ) = prediction_scores

        # Calculate the loss
        total_loss = None
        if (
            masked_lm_labels is not None
            and ent_masked_lm_labels is not None
            and prot_masked_lm_labels is not None
        ):
            loss_fct = nn.CrossEntropyLoss()
            # 1. Text-based MLM
            masked_lm_loss = loss_fct(
                token_prediction_scores.view(-1, self.config.vocab_size),
                masked_lm_labels.view(-1),
            )
            # 2. Entity-based masked "language" (entity) modeling
            ent_masked_lm_loss = loss_fct(
                entity_predictions_scores.view(-1, self.config.kg_vocab_size),
                ent_masked_lm_labels.view(-1),
            )
            # 3. Protein-based masked "language" (entity) modeling
            prot_masked_lm_loss = loss_fct(
                prot_predictions_scores.view(-1, self.config.prot_vocab_size),
                prot_masked_lm_labels.view(-1),
            )
            # Total loss = the sum of the individual training objective losses
            total_loss = masked_lm_loss + ent_masked_lm_loss + prot_masked_lm_loss

        if not return_dict:
            output = prediction_scores + outputs[2:]
            return ((total_loss,) + output) if total_loss is not None else output

        return BigBirdForPreTrainingOutputWithPooling(
            loss=total_loss,
            prediction_logits=prediction_scores,
            hidden_states=sequence_output,
            attentions=outputs.attentions,
            pooler_output=pooled_output,
        )
``` 

and I am using this model in a Trainer: 

```
    # Initialize the Trainer
    trainer = Trainer(
        model=stonkgs_model,
        args=training_args,
        train_dataset=pretraining_data,
    )
    train_result = trainer.train(resume_from_checkpoint=last_checkpoint)
```
This results in the following error: 

```
Traceback (most recent call last):
  File "stonkgs_pretraining.py", line 234, in <module>
    pretrain_stonkgs()
  File "/home/hbalabin/software/conda/envs/stonkgs/lib/python3.8/site-packages/click/core.py", line 829, in __call__
    return self.main(*args, **kwargs)
  File "/home/hbalabin/software/conda/envs/stonkgs/lib/python3.8/site-packages/click/core.py", line 782, in main
    rv = self.invoke(ctx)
  File "/home/hbalabin/software/conda/envs/stonkgs/lib/python3.8/site-packages/click/core.py", line 1066, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/home/hbalabin/software/conda/envs/stonkgs/lib/python3.8/site-packages/click/core.py", line 610, in invoke
    return callback(*args, **kwargs)
  File "stonkgs_pretraining.py", line 220, in pretrain_stonkgs
    train_result = trainer.train(resume_from_checkpoint=last_checkpoint)
  File "/home/hbalabin/software/conda/envs/stonkgs/lib/python3.8/site-packages/transformers/trainer.py", line 1284, in train
    tr_loss += self.training_step(model, inputs)
  File "/home/hbalabin/software/conda/envs/stonkgs/lib/python3.8/site-packages/transformers/trainer.py", line 1799, in training_step
    self.scaler.scale(loss).backward()
  File "/home/hbalabin/software/conda/envs/stonkgs/lib/python3.8/site-packages/torch/_tensor.py", line 255, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/home/hbalabin/software/conda/envs/stonkgs/lib/python3.8/site-packages/torch/autograd/__init__.py", line 147, in backward
    Variable._execution_engine.run_backward(
RuntimeError: Trying to backward through the graph a second time (or directly access saved variables after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved variables after calling backward.
``` 

## Expected behavior

My guess is that this error message is caused by loss.backward() being called multiple times (maybe there is something wrong with the gradient accumulation procedure). The expected behavior would be that the backward procedure is only applied once to the graph/that loss.backward() being called once per training step. 

Has anyone seen this error before? Thanks a lot in advance!! :hugs: 


Comments:
You should debug your custom model on your own training loop first (pass a batch, compute the loss and call backward), as it's more likely to come from your custom model than the `Trainer`.

Makes sense, I'll try that out and see if I can get any insights from that. Thanks! :smile: 

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.