gradient checkpointing disables requires_grad when freezing part of models (fix with use_reentrant=False)
Author: PaulLerner
Pull request: no

### System Info

- `transformers` version: 4.22.2
- Platform: Linux-4.18.0-372.36.1.el8_6.x86_64-x86_64-with-glibc2.28
- Python version: 3.10.4
- Huggingface_hub version: 0.10.0
- PyTorch version (GPU?): 1.12.1 (False)
- 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 but not relevant
- Using distributed or parallel set-up in script?: no but not relevant

### Who can help?

trainer/PyTorch: @sgugger

### Information

- [ ] The official example scripts
- [X] 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

When freezing the first layers of a model (e.g. including embeddings) *and* using gradient-checkpointing, **all** gradient calculation will be disabled, i.e. the output will have `requires_grad==False`. See https://discuss.pytorch.org/t/checkpoint-with-no-grad-requiring-inputs-problem/19117/7 for explanations. I believe this is the issue encountered by @entslscheia in https://github.com/huggingface/transformers/issues/16276 (apparently unsolved).

Here is a sample code to reproduce with `BertModel`, but I think that gradient-checkpointing is implemented like this everywhere in the library:
```py
In [1]: from transformers import BertTokenizer, BertModel
In [2]: model = BertModel.from_pretrained('../models/bert-base-uncased/')                                                                  
In [3]: tokenizer = BertTokenizer.from_pretrained('../models/bert-base-uncased/')                                                          
In [4]: inputs = tokenizer(['foo', 'bar'], return_tensors='pt')
# enable gradient checkpointing
In [5]: model.encoder.gradient_checkpointing = True
# freezing the first input layers (here the embeddings)
In [8]: for p in model.embeddings.parameters():
   ...:     p.requires_grad=False
In [9]: output = model(**inputs)                     
# expected True                                                                                      
In [12]: output.last_hidden_state.requires_grad
Out[12]: False
# note that all weights of the model have requires_grad==True except for the embeddings
In [15]: model.encoder.layer[0].output.dense.weight.requires_grad
Out[15]: True
```
As myself, you might find out about this while training a model, because obviously if you pass `None` gradients to an optimizer, it will not be happy. You might encounter the following warning: `/gpfswork/rech/fih/usl47jg/miniconda3/envs/datasets/lib/python3.10/site-packages/torch/utils/checkpoint.py:25: UserWarning: None of the inputs have requires_grad=True. Gradients will be None` and then the following error: `RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn`.


### Expected behavior

A quick fix for me was to set `use_reentrant=False` when calling `torch.utils.checkpoint.checkpoint` (e.g. in https://github.com/huggingface/transformers/blob/main/src/transformers/models/bert/modeling_bert.py#L600). Note that this will be the future default in "future versions of PyTorch" ([according to the doc](https://pytorch.org/docs/stable/checkpoint.html) without further precision).

Comments:
Thanks for flagging the issue. We don't offer APIs to freeze the model in the Trainer (as it has never shown better results for fine-tuning, quite the opposite), so we can leave this issue to show how to solve the problem, but we won't really incorporate it in Transformers.

I think a lot of people use `transformers` models without using the Trainer (so with their own training script or, e.g. pytorch lightning that *does* provide an API to freeze models), but it’s your call :)

in case someone needs quick fix for their custom training pipelines
For inspiration:
```python
    if gradient_checkpointing_enabled:
        from functools import partial
        notfailing_checkpoint = partial(torch.utils.checkpoint.checkpoint, use_reentrant=False)
        torch.utils.checkpoint.checkpoint = notfailing_checkpoint
        model.gradient_checkpointing_enable()
```

qq: this mitigation is tested on only 1 GPU or multiple GPUs? On multiple GPUs, I saw this error: 
AssertionError: Expects storage to be allocated

I only tested with 1 GPU

> I only tested with 1 GPU

as I expected. Let's wait for some updates from the pytorch side.

Just hit this today... a simple and efficient way to fine tune LLMs is to just train some layers, and of course checkpointing would help a lot here, so IMO would be great to be able to specify `use_reentrant` in Transformers!

Hi @harpone 
With https://github.com/huggingface/transformers/pull/27020 being merged, you can do
```python
model.enable_gradient_checkpointing(gradient_checkpointing_kwargs={"use_reentrant": False})
```
To enable `use_reentrant=False`.