Pytorch - Lazy initialization of models

[patrickvonplaten]

🚨🚨🚨 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:

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:

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

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 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

[sgugger]

Thanks for tackling this! I have a few comments on the design, let's continue the discussion on each one!

[stas00]

So glad we finally are trying to sort it out!

If we are still uncertain about some edge cases perhaps we can also empower the user to have a way to fully init weights and bypass this optimization

[LysandreJik]

Cool, great to see such a great speedup! I'd rather be on the more cautious side here and make sure that each weight is correctly initialized, even if it's done twice, rather than have some weights which are not initialized.

[stas00]

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

[john-heyer]

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!

[patrickvonplaten]

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?

[patrickvonplaten]

@john-heyer - I answered in-detail here: #17370

[john-heyer]

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.

[vadimkantorov]

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)`

[patrickvonplaten]

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

[vadimkantorov]

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