[`GPTNeo`] Fix gradient checkpointing bug
Author: younesbelkada
Pull request: yes

# What does this PR do?
 
This PR fixes a small bug that a user can encounter while using `generate` and models that use `gradient_checkpointing` . In the context of `trl` we call `generate` on the active model that uses `gradient_checkpointing` to save memory. 

Currently this snippet fails on `main`: 

```python
import torch

from transformers import AutoModelForCausalLM

# Now let's build the model, the reference model, and the tokenizer.
model = AutoModelForCausalLM.from_pretrained("edbeeching/gpt-neo-125M-imdb", device_map="auto")
model.train()
model.gradient_checkpointing_enable()
gen = model.generate(input_ids=torch.LongTensor([0,1,2,3]).unsqueeze(0)) 
```

This is because `gradient_checkpointing_enable` and `use_cache` are not compatible, and calling `generate` uses `use_cache` by default. Though there is a check to force `use_cache` to be set to `False`, this check is not enough since the `present` tuple will be still tried to be populated but remains empty since it will be intialized with an empty tuple, Hence leading to a confusing error.
IMO the fix should be to force-set `presents` tuple to `None` if the model is using `gradient_checkpointing`

cc @sgugger @amyeroberts 

Comments:
_The documentation is not available anymore as the PR was closed or merged._

cc @gante 

Change looks reasonable to me. Could you provide links to the check you refer to where `use_cache` is set to False and also add the error output so this solution can be found easily if encountered again please? 

Sure, here is the traceback: 
```bash
│ /home/younes_huggingface_co/code/transformers/src/transformers/generation/utils.py:1402 in       │
│ generate                                                                                         │
│                                                                                                  │
│   1399 │   │   │   │   )                                                                         │
│   1400 │   │   │                                                                                 │
│   1401 │   │   │   # 11. run greedy search                                                       │
│ ❱ 1402 │   │   │   return self.greedy_search(                                                    │
│   1403 │   │   │   │   input_ids,                                                                │
│   1404 │   │   │   │   logits_processor=logits_processor,                                        │
│   1405 │   │   │   │   stopping_criteria=stopping_criteria,                                      │
│                                                                                                  │
│ /home/younes_huggingface_co/code/transformers/src/transformers/generation/utils.py:2197 in       │
│ greedy_search                                                                                    │
│                                                                                                  │
│   2194 │   │   │   model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)  │
│   2195 │   │   │                                                                                 │
│   2196 │   │   │   # forward pass to get next token                                              │
│ ❱ 2197 │   │   │   outputs = self(                                                               │
│   2198 │   │   │   │   **model_inputs,                                                           │
│   2199 │   │   │   │   return_dict=True,                                                         │
│   2200 │   │   │   │   output_attentions=output_attentions,                                      │
│                                                                                                  │
│ /home/younes_huggingface_co/miniconda3/envs/fix-test/lib/python3.9/site-packages/torch/nn/module │
│ s/module.py:1194 in _call_impl                                                                   │
│                                                                                                  │
│   1191 │   │   # this function, and just call forward.                                           │
│   1192 │   │   if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks o  │
│   1193 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks):                   │
│ ❱ 1194 │   │   │   return forward_call(*input, **kwargs)                                         │
│   1195 │   │   # Do not call functions when jit is used                                          │
│   1196 │   │   full_backward_hooks, non_full_backward_hooks = [], []                             │
│   1197 │   │   if self._backward_hooks or _global_backward_hooks:                                │
│                                                                                                  │
│ /home/younes_huggingface_co/miniconda3/envs/fix-test/lib/python3.9/site-packages/accelerate/hook │
│ s.py:158 in new_forward                                                                          │
│                                                                                                  │
│   155 │   │   │   with torch.no_grad():                                                          │
│   156 │   │   │   │   output = old_forward(*args, **kwargs)                                      │
│   157 │   │   else:                                                                              │
│ ❱ 158 │   │   │   output = old_forward(*args, **kwargs)                                          │
│   159 │   │   return module._hf_hook.post_forward(module, output)                                │
│   160 │                                                                                          │
│   161 │   module.forward = new_forward                                                           │
│                                                                                                  │
│ /home/younes_huggingface_co/code/transformers/src/transformers/models/gpt_neo/modeling_gpt_neo.p │
│ y:739 in forward                                                                                 │
│                                                                                                  │
│   736 │   │   """                                                                                │
│   737 │   │   return_dict = return_dict if return_dict is not None else self.config.use_return   │
│   738 │   │                                                                                      │
│ ❱ 739 │   │   transformer_outputs = self.transformer(                                            │
│   740 │   │   │   input_ids,                                                                     │
│   741 │   │   │   past_key_values=past_key_values,                                               │
│   742 │   │   │   attention_mask=attention_mask,                                                 │
│                                                                                                  │
│ /home/younes_huggingface_co/miniconda3/envs/fix-test/lib/python3.9/site-packages/torch/nn/module │
│ s/module.py:1194 in _call_impl                                                                   │
│                                                                                                  │
│   1191 │   │   # this function, and just call forward.                                           │
│   1192 │   │   if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks o  │
│   1193 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks):                   │
│ ❱ 1194 │   │   │   return forward_call(*input, **kwargs)                                         │
│   1195 │   │   # Do not call functions when jit is used                                          │
│   1196 │   │   full_backward_hooks, non_full_backward_hooks = [], []                             │
│   1197 │   │   if self._backward_hooks or _global_backward_hooks:                                │
│                                                                                                  │
│ /home/younes_huggingface_co/code/transformers/src/transformers/models/gpt_neo/modeling_gpt_neo.p │
│ y:545 in forward                                                                                 │
│                                                                                                  │
│   542 │   │   │   past_length = 0                                                                │
│   543 │   │   │   past_key_values = tuple([None] * len(self.h))                                  │
│   544 │   │   else:                                                                              │
│ ❱ 545 │   │   │   past_length = past_key_values[0][0].size(-2)                                   │
│   546 │   │                                                                                      │
│   547 │   │   if position_ids is None:                                                           │
│   548 │   │   │   position_ids = torch.arange(past_length, input_shape[-1] + past_length, dtyp   │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
IndexError: tuple index out of range
```

And `use_cache` is force-set to False [here](https://github.com/huggingface/transformers/blob/aff87da15b04b260c6057dd47a70376f2b2386f3/src/transformers/models/gpt_neox/modeling_gpt_neox.py#L515) but `presents` are already initialized [here](https://github.com/huggingface/transformers/blob/aff87da15b04b260c6057dd47a70376f2b2386f3/src/transformers/models/gpt_neox/modeling_gpt_neox.py#L503), hence returned [here](https://github.com/huggingface/transformers/blob/aff87da15b04b260c6057dd47a70376f2b2386f3/src/transformers/models/gpt_neox/modeling_gpt_neox.py#L555)


The reason for the change sounds good 👍 

Two comments, though:
1. We can solve this problem without adding more logic -- e.g. moving `if` that sets `use_cache=False` up in the class, before `presents` is initialized. Less logic = simpler code = fewer bugs = 💛  
2. This is a widespread pattern in generate-compatible models. If we make this change, the least we should do is open an issue with the `Good First Issue` label that tracks which models have already received the change!

Thanks a lot for these suggestions @gante! This makes a lot of sense, 
I have adapted the code accordingly and drafted a Good First issue that we can edit once we figure out if the bug persists for other models

I'm getting same error in t5, is this the same reason ?

File /usr/local/lib/python3.8/dist-packages/transformers/models/t5/modeling_t5.py:506, in T5Attention.forward.<locals>.project(hidden_states, proj_layer, key_value_states, past_key_value)
    502 if past_key_value is not None:
    503     if key_value_states is None:
    504         # self-attn
    505         # (batch_size, n_heads, key_length, dim_per_head)
--> 506         hidden_states = torch.cat([past_key_value, hidden_states], dim=2)
    507     elif past_key_value.shape[2] != key_value_states.shape[1]:
    508         # checking that the `sequence_length` of the `past_key_value` is the same as
    509         # the provided `key_value_states` to support prefix tuning
    510         # cross-attn
    511         # (batch_size, n_heads, seq_length, dim_per_head)
    512         hidden_states = shape(proj_layer(key_value_states))

RuntimeError: Sizes of tensors must match except in dimension 2. Expected size 16 but got size 32 for tensor number 1 in the list.

Not sure in this case, can you share a reproducible script?