Pytorch - Lazy initialization of models
Author: patrickvonplaten
Pull request: yes

🚨🚨🚨 **Breaking seeded model initialization** 🚨🚨🚨

As explained below this PR breaks seeded model initialization by default. To ensure the exact same model initialization as before, use:

```python
torch.manual_seed(seed)
model = BertForSequenceClassification.from_pretrained("bert-base-cased", _fast_init=False)
```

# What does this PR do?

This PR implements fast initializing by only initializing weights that need to be initialized. For every model two aggressive tests are added to make sure that the new "fast" initialization initializes the weights according to the exact same distributions as the previous init scheme. 

IMO, it is not possible to ensure that:

```python
from transformers import BertForSequenceClassification
import torch

torch.manual_seed(0)
model = BertForSequenceClassification.from_pretrained("bert-base-cased", _fast_init=False)  # this randomely inits the lm_head layer
```

yields the same results as the new "fast" init

```python
torch.manual_seed(0)
model = BertForSequenceClassification.from_pretrained("bert-base-cased")  # this randomely inits the lm_head layer
```

since in the first case all layers are initialized so that the random number generator is called much more often, thus making it impossible to ensure identical weights initialization => compare to [this](https://discuss.pytorch.org/t/does-pytorch-change-its-internal-seed-during-training/46505/4) post to better understand why this is probably not possible.

This became obvious in this PR since I had to change the random_seed for running the `run_ner.py` examples test to make it pass.

I guess it is therefore better to stick to "initializing all weights" for now and only when changing to version 5.0 making this breaking change. => Guess we should discuss this @sgugger @LysandreJik @stas00 


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

Fixes: #9205

Comments:
Not to impede this PR, just wanted to add that this comment might be of importance for future work in this area: https://github.com/pytorch/pytorch/issues/29523#issuecomment-831390099

And the whole issue: https://github.com/pytorch/pytorch/issues/29523

I wonder if our side needs/should support that `reset_parameters` feature in our models. Given the last comment https://github.com/pytorch/pytorch/issues/29523#issuecomment-831435863 it's unclear where it's standing. So perhaps this is something to revisit later when the dust settles on the pytorch side.

FYI, pytorch has just added `torch.nn.utils.skip_init()` to handle similar situations:
https://pytorch.org/tutorials/prototype/skip_param_init.html
it should appear probably around pt-1.9.1.

Note that `torch.nn.utils.skip_init()` is even more efficient as it doesn't allocate any storage at all! So there is not even an overhead of creating any weights until they are loaded from state_dict. https://pytorch.org/tutorials/prototype/skip_param_init.html#implementation-details


Hey all - this introduced a pretty nasty bug for us, that took a while to figure out. Here's a case where this initialization of the `missing_keys`, which normally shouldn't matter, broke our model after version 4.6. We have a custom model that optionally initializes some of its weights from a separate module. Calling `from_pretrained` and passing this separate module used to "work", but now those weights are overwritten. See MWE:
```
from torch.nn import Linear
from transformers import BertModel


class MyCustomModel(BertModel):
    def __init__(self, config, custom_layer=None):
        super().__init__(config)
        if custom_layer is not None:
            self.custom_layer = custom_layer
        else:
            self.custom_layer = Linear(1024, 1024)


if __name__ == "__main__":
    import transformers
    print(transformers.__version__)
    layer = Linear(1024, 1024)
    print(layer.weight.sum())
    custom_model = MyCustomModel.from_pretrained('bert-base-uncased', custom_layer=layer)
    # used to be the same as the layer above, but it is "re-initialized" in the from_pretrained method
    print(custom_model.custom_layer.weight.sum())
```
Result:
```
4.11.3
tensor(5.9874, grad_fn=<SumBackward0>)
Some weights of the model checkpoint at bert-base-uncased were not used when initializing MyCustomModel: ['cls.seq_relationship.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.bias', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.decoder.weight', 'cls.seq_relationship.bias']
- This IS expected if you are initializing MyCustomModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
- This IS NOT expected if you are initializing MyCustomModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
Some weights of MyCustomModel were not initialized from the model checkpoint at bert-base-uncased and are newly initialized: ['bert.custom_layer.bias', 'bert.custom_layer.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
tensor(17.4651, grad_fn=<SumBackward0>)
```

This used to "just work" in version < 4.6. Perhaps we were relying on an unintended "feature". Setting `_fast_init=False` *does* fix things, but it's a bit hacky as it only applies to the initialization of this custom module that is called upstream in our service. Additionally, we don't know what happens if we'll need to rely on this feature in the future, but it goes away.

Can you comment on this? Thanks!

Hey @john-heyer,

Thanks for the feedback. We indeed didn't take account the effect this would have on custom models that inherit from `transformers` models like BERT. Just to understand better, the problem was that before 4.6, the `custom_layer` was not (re-)initialized when calling `MyCustomModel.from_pretrained(...)` - however after 4.6 it was initialized twice once before calling `from_pretrained(...)` and once after calling it? 

@john-heyer - I answered in-detail here: https://github.com/huggingface/transformers/issues/17370

thanks @patrickvonplaten - yes, that is correct - before 4.6 `custom_layer` was not re-initialized, and now it is! Thanks for opening the other thread. I'll follow there.

Can this be used to load pretrained params directly to the GPU models without keeping the full copy of params on CPU?

E.g. https://github.com/huggingface/diffusers has `unet = UNetModel.from_pretrained("fusing/ddpm-lsun-church").to(torch_device)` while theoretically one could have unet = UNetModel.from_pretrained("fusing/ddpm-lsun-church"m device = torch_device)`

Hey @vadimkantorov, could you maybe open an issue under `diffusers` instead? :-)

Is transfomers following a different design? I assumed diffusers just copied the original design from transformers