Potential bug in modeling_utils.py
Author: NielsRogge
Pull request: no

### System Info

Transformers main branch.



### Who can help?

@LysandreJik @sgugger 

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

There currently seems to be a potential bug in modeling_utils.py, revealed by the failure of 2 tests defined in `test_modeling_common.py`, namely [test_correct_missing_keys](https://github.com/huggingface/transformers/blob/243439a8271137aa290d7546e5704feeaa0cd1e5/tests/test_modeling_common.py#L1448) and [test_save_load_fast_init_from_base](https://github.com/huggingface/transformers/blob/243439a8271137aa290d7546e5704feeaa0cd1e5/tests/test_modeling_common.py#L313).

The issue occurs when a head model (like `xxxForSequenceClassification`) defines a parameter that has the same name as one in the base model (`xxxModel`). Let's say the base model defines a `self.layernorm` attribute/parameter, and the head model also defines a `self.layernorm`.

You can reproduce the error by cloning [this branch](https://github.com/NielsRogge/transformers/tree/add_ast_bug) of mine, then run the following tests:

```
pytest tests/models/audio_spectogram_transformer/test_modeling_audio_spectogram_transformer.py::AudioSpectogramTransformerModelTest
```

In that case, both tests fail with the following error:

```
(...)
# Some models may have keys that are not in the state by design, removing them before needlessly warning
        # the user.
        if cls._keys_to_ignore_on_load_missing is not None:
            for pat in cls._keys_to_ignore_on_load_missing:
                missing_keys = [k for k in missing_keys if re.search(pat, k) is None]
    
        if cls._keys_to_ignore_on_load_unexpected is not None:
            for pat in cls._keys_to_ignore_on_load_unexpected:
                unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None]
    
        # retrieve weights on meta device and put them back on CPU.
        # This is not ideal in terms of memory, but if we don't do that not, we can't initialize them in the next step
        if low_cpu_mem_usage:
            for key in missing_keys:
                if key.startswith(prefix):
                    key = ".".join(key.split(".")[1:])
                param = model_state_dict[key]
                if param.device == torch.device("meta"):
                    if not load_in_8bit:
                        set_module_tensor_to_device(model, key, "cpu", torch.empty(*param.size(), dtype=dtype))
                    else:
                        set_module_8bit_tensor_to_device(model, key, "cpu", torch.empty(*param.size(), dtype=dtype))
    
        # retrieve unintialized modules and initialize before maybe overriding that with the pretrained weights.
        if _fast_init:
            uninitialized_modules = model.retrieve_modules_from_names(
                missing_keys, add_prefix=add_prefix_to_model, remove_prefix=remove_prefix_from_model
            )
            for module in uninitialized_modules:
                model._init_weights(module)
    
        # Make sure we are able to load base models as well as derived models (with heads)
        start_prefix = ""
        model_to_load = model
        if len(cls.base_model_prefix) > 0 and not hasattr(model, cls.base_model_prefix) and has_prefix_module:
            start_prefix = cls.base_model_prefix + "."
        if len(cls.base_model_prefix) > 0 and hasattr(model, cls.base_model_prefix) and not has_prefix_module:
            model_to_load = getattr(model, cls.base_model_prefix)
            if any(key in expected_keys_not_prefixed for key in loaded_keys):
>               raise ValueError(
                    "The state dictionary of the model you are trying to load is corrupted. Are you sure it was "
                    "properly saved?"
                )
E               ValueError: The state dictionary of the model you are trying to load is corrupted. Are you sure it was properly saved?
```

However, when simply renaming `self.layernorm` to `self.layer_norm` in the head model, both tests pass.

### Expected behavior

Normally, this should work without any error. I think the reason we haven't encountered this issue yet is simply because the case where a head model defines a parameter that has the same name as one defined in the base model is quite rare. However, this should still work as expected, as in this case for instance the layernorm of the base Transformer will have `audio_spectogram_transformer.layernorm` as name and the layernorm of the head model simply `layernorm`.

Unless I'm missing something here ;) happy to discuss.

Comments:
Will have a look this morning. Sounds like a bug indeed!

@sgugger thanks for fixing, however I'm still encountering an issue that is probably related to this. 

Specifically, the parameter whose name is the same between a base model and a head model (`self.layernorm` in my case) makes the [test_save_load_fast_init_from_base](https://github.com/huggingface/transformers/blob/243439a8271137aa290d7546e5704feeaa0cd1e5/tests/test_modeling_common.py#L313) test fail.

It can be reproduced as follows:

```
RUN_SLOW=yes pytest tests/models/audio_spectrogram_transformer/test_modeling_audio_spectrogram_transformer.py::AudioSpectrogramTransformerModelTest::test_save_load_fast_init_from_base
```
This might also be related to the fast init mechanism itself, which doesn't seem to support parameters which have the same name between the base and head model. Should I just skip the test?


No, as this means parameters won't be properly initialized. I won't have any bandwidth to fix this in the near future so someone else will have to fix it.

I had a very quick look into it and I don't see any easy fix -> so for now I would advise to use a different name for the weights in the head (like `final_layernorm` maybe?)

Ok, I'll do that, thanks for looking into it

Same error here: ValueError: The state dictionary of the model you are trying to load is corrupted. Are you sure it was properly saved?