Add `keep_in_fp32_modules` support
Author: younesbelkada
Pull request: yes

# What does this PR do?

This PR partially addresses #20287 - although half-precision and int8 conversion work extremely well for most of the models, for some architectures (e.g. T5) the casting leads in a drastic performance degradation. 

This can be fixed by manually force-casting some modules in `float32`. For FLAN-T5, @larsmennen and @navjotts have found out that keeping only these weights in `fp32` enables to run largest models in fp16 or int8 with no performance degradation.

This PR introduces a new utils in `from_pretrained` method, termed as `keep_in_fp32_modules` that partially addresses this issue.

How this util works? For T5:
```
from transformers import T5ForConditionalGeneration

model = T5ForConditionalGeneration.from_pretrained("t5-small", torch_dtype=torch.float16, keep_in_fp32_modules=["wo"])
print(model.decoder.block[0].layer[2].DenseReluDense.wo.weight.dtype)
>>> torch.float32
```

When using `keep_in_fp32_modules` , `low_cpu_mem_usage` needs to be force-set to `True`. This is because if `low_cpu_mem_usage=False`, it is the function from pytorch `_load_from_state_dict` that is called under the hood on each sub-module. This function calls `copy_` from Pytorch which seems to keep the tensor in its native `dtype` regardless the `dtype` of the input module

```
import torch

param = torch.Tensor([0.1, 0.2, 0.3]).to(torch.float16)
to_copy_param = torch.Tensor([0.2, 0.1, 0.3]).to(torch.float32)

param.copy_(to_copy_param)
print(param.dtype)
>>> torch.float16
```

Keeping this as a draft for now as this util needs to be manually patched with fixes such as https://github.com/huggingface/transformers/issues/20287#issuecomment-1342219429 , otherwise users will encounter issues about incompatible `dtype` between input and weights

cc @sgugger 

Comments:
What about adding hooks on each converted module, that will take care of converting the input / output to the correct `dtype` ?

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

As suggested in #20287 / model loaded in bfloat16 should keep their weights in `bfloat16` and not cast them in `fp32`. This is addressed in  e3498da

Thanks @younesbelkada and @sgugger !! Tested this locally; can confirm this works with patch 1&2 from https://github.com/huggingface/transformers/issues/20287#issuecomment-1342219429

The only problem I encountered is that in: 

https://github.com/huggingface/transformers/blob/0f75387d633619b8c0cf6955b07dac54d2e17473/src/transformers/modeling_utils.py#L2326

You get an error as `keep_in_fp32_modules` is an unexpected keyword to the underlying model class (locally i just added it quickly to test). Do you want to add this in so people can use it in their model class to determine where to apply patches like 1&2? Or alternatively don't pass it on and then people can just query the dtype.

Thanks so much @larsmennen for confirming that the tests pass! We should be close merging this 💪 
I think that your failing test should be fixed with my latest commit ( cb89c42 ) but I am not sure, could you try again with the latest commit? 🙏 

hmm that doesn't fix it. I think you just need to pop the argument from model_kwargs, otherwise it gets passed to the underlying model (i'm assuming you don't want that? but cmiiw)

I.e. after

https://github.com/huggingface/transformers/blob/7d47df2e52c6b55d8f19e5b2bd8b5e472a4f0a82/src/transformers/modeling_utils.py#L1981

if you add

```py
keep_in_fp32_modules = kwargs.pop("keep_in_fp32_modules", None)
```

I tested w/ that modification on top of [7d47df2](https://github.com/huggingface/transformers/pull/20683/commits/7d47df2e52c6b55d8f19e5b2bd8b5e472a4f0a82) and that works! Thanks for the quick action @younesbelkada ! 🙏 

@larsmennen how are you loading your model ? The description above is slightly misleading as initially the plan was to add a kwarg when loading the model as follows:
```
from transformers import T5ForConditionalGeneration

model = T5ForConditionalGeneration.from_pretrained("t5-small", torch_dtype=torch.float16, keep_in_fp32_modules=["wo"])
```
but now this is not needed, you should just load your model like:
```
from transformers import T5ForConditionalGeneration

model = T5ForConditionalGeneration.from_pretrained("t5-small", device_map="auto", load_in_8bit=True])
```

@younesbelkada ah i see! I was passing the kwarg yes, so that explains.

@larsmennen this PR will be merged as soon as all the tests will be green ! 
Would you mind opening a PR addressing your suggestions (patch 1 & 2 from the discussion at #20287 )? 

All slow tests from T5 (and BLOOM just in case we didn't break anything else) pass 🟢 
Merging once the CI tests are green

I tried this one, latest version of transformers (27.4), cuda 10.2 and I get this error:

`model1a_CPU = T5ForConditionalGeneration.from_pretrained(model_path, low_cpu_mem_usage=True,torch_dtype=torch.float16, keep_in_fp32_modules=["wo"]).to("cuda")               
TypeError: __init__() got an unexpected keyword argument 'keep_in_fp32_modules'`




`keep_in_fp32_modules` is not an argument you can pass to `from_pretrained`, this is all done internally.

You need to do somthing like:

```python
from transformers import T5ForConditionalGeneration

T5ForConditionalGeneration._keep_in_fp32_modules = ["wo"]

# your code here
```

Except this is already done for T5 ;-)