FlaxRobertaModel.from_pretrained does not load weights correctly
Author: bminixhofer
Pull request: no

## Environment info
- `transformers` version: 4.8.2
- Platform: Linux-5.4.104+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.10
- PyTorch version (GPU?): 1.9.0+cu102 (False)
- Tensorflow version (GPU?): 2.5.0 (False)
- Flax version (CPU?/GPU?/TPU?): 0.3.4 (cpu)
- Jax version: 0.2.13
- JaxLib version: 0.1.66
- Using GPU in script?: no
- Using distributed or parallel set-up in script?: no

## Information

`from_pretrained` does not work correctly for the `roberta-base` Flax model. I have not checked if this also affects other models.

## To reproduce

Steps to reproduce the behavior:

```python
from transformers import FlaxRobertaModel

model = FlaxRobertaModel.from_pretrained("roberta-base")
```

prints:

```
Some weights of the model checkpoint at roberta-base were not used when initializing FlaxRobertaModel: {('encoder', 'layer', '9', 'attention', 'layer_norm', 'beta'), ('encoder', 'layer', '0', 'attention', 'self', 'out', 'kernel'), ('encoder', 'layer', '2', 'output', 'layer_norm', 'beta'), ('embeddings', 'word_embeddings', 'weight'), ('encoder', 'layer', '4', 'attention', 'layer_norm', 'gamma'), ('encoder', 'layer', '4', 'output', 'layer_norm', 'gamma'), ('encoder', 'layer', '8', 'output', 'layer_norm', 'beta'), ('encoder', 'layer', '10', 'attention', 'self', 'out', 'bias'), ('encoder', 'layer', '1', 'attention', 'layer_norm', 'beta'), ('encoder', 'layer', '11', 'attention', 'self', 'out', 'kernel'), ('encoder', 'layer', '10', 'attention', 'layer_norm', 'beta'), ('encoder', 'layer', '6', 'attention', 'self', 'out', 'bias'), ('encoder', 'layer', '7', 'attention', 'self', 'out', 'kernel'), ('encoder', 'layer', '5', 'output', 'layer_norm', 'beta'), ('layer_norm', 'weight'), ('encoder', 'layer', '7', 'output', 'layer_norm', 'gamma'), ('encoder', 'layer', '11', 'output', 'layer_norm', 'gamma'), ('encoder', 'layer', '4', 'attention', 'self', 'out', 'bias'), ('encoder', 'layer', '0', 'attention', 'self', 'out', 'bias'), ('encoder', 'layer', '1', 'output', 'layer_norm', 'beta'), ('encoder', 'layer', '5', 'attention', 'self', 'out', 'bias'), ('encoder', 'layer', '5', 'attention', 'layer_norm', 'beta'), ('encoder', 'layer', '1', 'attention', 'self', 'out', 'kernel'), ('encoder', 'layer', '3', 'output', 'layer_norm', 'gamma'), ('encoder', 'layer', '10', 'output', 'layer_norm', 'gamma'), ('encoder', 'layer', '0', 'attention', 'layer_norm', 'beta'), ('encoder', 'layer', '3', 'attention', 'self', 'out', 'kernel'), ('encoder', 'layer', '6', 'output', 'layer_norm', 'beta'), ('encoder', 'layer', '8', 'attention', 'self', 'out', 'bias'), ('encoder', 'layer', '0', 'output', 'layer_norm', 'gamma'), ('encoder', 'layer', '7', 'attention', 'layer_norm', 'gamma'), ('encoder', 'layer', '8', 'attention', 'layer_norm', 'gamma'), ('encoder', 'layer', '1', 'attention', 'layer_norm', 'gamma'), ('encoder', 'layer', '2', 'attention', 'layer_norm', 'gamma'), ('encoder', 'layer', '9', 'attention', 'layer_norm', 'gamma'), ('decoder', 'weight'), ('encoder', 'layer', '9', 'attention', 'self', 'out', 'kernel'), ('encoder', 'layer', '2', 'attention', 'self', 'out', 'kernel'), ('encoder', 'layer', '11', 'attention', 'layer_norm', 'gamma'), ('encoder', 'layer', '3', 'attention', 'layer_norm', 'gamma'), ('encoder', 'layer', '4', 'attention', 'self', 'out', 'kernel'), ('encoder', 'layer', '7', 'attention', 'self', 'out', 'bias'), ('encoder', 'layer', '11', 'output', 'layer_norm', 'beta'), ('encoder', 'layer', '9', 'output', 'layer_norm', 'beta'), ('encoder', 'layer', '2', 'output', 'layer_norm', 'gamma'), ('encoder', 'layer', '6', 'attention', 'self', 'out', 'kernel'), ('encoder', 'layer', '11', 'attention', 'self', 'out', 'bias'), ('encoder', 'layer', '6', 'attention', 'layer_norm', 'beta'), ('encoder', 'layer', '4', 'output', 'layer_norm', 'beta'), ('encoder', 'layer', '10', 'attention', 'self', 'out', 'kernel'), ('encoder', 'layer', '5', 'attention', 'layer_norm', 'gamma'), ('encoder', 'layer', '5', 'output', 'layer_norm', 'gamma'), ('embeddings', 'token_type_embeddings', 'weight'), ('layer_norm', 'bias'), ('encoder', 'layer', '10', 'attention', 'layer_norm', 'gamma'), ('embeddings', 'layer_norm', 'gamma'), ('encoder', 'layer', '1', 'output', 'layer_norm', 'gamma'), ('encoder', 'layer', '6', 'output', 'layer_norm', 'gamma'), ('encoder', 'layer', '4', 'attention', 'layer_norm', 'beta'), ('encoder', 'layer', '2', 'attention', 'self', 'out', 'bias'), ('encoder', 'layer', '8', 'output', 'layer_norm', 'gamma'), ('encoder', 'layer', '9', 'attention', 'self', 'out', 'bias'), ('encoder', 'layer', '3', 'output', 'layer_norm', 'beta'), ('dense', 'kernel'), ('encoder', 'layer', '8', 'attention', 'layer_norm', 'beta'), ('encoder', 'layer', '10', 'output', 'layer_norm', 'beta'), ('encoder', 'layer', '8', 'attention', 'self', 'out', 'kernel'), ('bias',), ('dense', 'bias'), ('encoder', 'layer', '3', 'attention', 'self', 'out', 'bias'), ('encoder', 'layer', '7', 'output', 'layer_norm', 'beta'), ('encoder', 'layer', '9', 'output', 'layer_norm', 'gamma'), ('encoder', 'layer', '0', 'output', 'layer_norm', 'beta'), ('embeddings', 'layer_norm', 'beta'), ('encoder', 'layer', '0', 'attention', 'layer_norm', 'gamma'), ('encoder', 'layer', '2', 'attention', 'layer_norm', 'beta'), ('encoder', 'layer', '6', 'attention', 'layer_norm', 'gamma'), ('encoder', 'layer', '1', 'attention', 'self', 'out', 'bias'), ('encoder', 'layer', '11', 'attention', 'layer_norm', 'beta'), ('embeddings', 'position_embeddings', 'weight'), ('encoder', 'layer', '5', 'attention', 'self', 'out', 'kernel'), ('encoder', 'layer', '3', 'attention', 'layer_norm', 'beta'), ('encoder', 'layer', '7', 'attention', 'layer_norm', 'beta')}
- This IS expected if you are initializing FlaxRobertaModel 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 FlaxRobertaModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
Some weights of FlaxRobertaModel were not initialized from the model checkpoint at roberta-base and are newly initialized: {('encoder', 'layer', '6', 'attention', 'output', 'dense', 'bias'), ('encoder', 'layer', '3', 'output', 'LayerNorm', 'bias'), ('encoder', 'layer', '8', 'attention', 'output', 'LayerNorm', 'bias'), ('encoder', 'layer', '8', 'attention', 'output', 'LayerNorm', 'scale'), ('encoder', 'layer', '4', 'attention', 'output', 'LayerNorm', 'scale'), ('encoder', 'layer', '9', 'output', 'LayerNorm', 'bias'), ('encoder', 'layer', '9', 'output', 'LayerNorm', 'scale'), ('encoder', 'layer', '9', 'attention', 'output', 'dense', 'bias'), ('encoder', 'layer', '7', 'output', 'LayerNorm', 'bias'), ('encoder', 'layer', '10', 'output', 'LayerNorm', 'scale'), ('encoder', 'layer', '9', 'attention', 'output', 'LayerNorm', 'scale'), ('embeddings', 'position_embeddings', 'embedding'), ('encoder', 'layer', '3', 'output', 'LayerNorm', 'scale'), ('encoder', 'layer', '2', 'output', 'LayerNorm', 'scale'), ('encoder', 'layer', '1', 'output', 'LayerNorm', 'bias'), ('encoder', 'layer', '5', 'attention', 'output', 'dense', 'bias'), ('encoder', 'layer', '5', 'attention', 'output', 'LayerNorm', 'bias'), ('encoder', 'layer', '5', 'attention', 'output', 'LayerNorm', 'scale'), ('encoder', 'layer', '0', 'attention', 'output', 'LayerNorm', 'bias'), ('encoder', 'layer', '6', 'output', 'LayerNorm', 'bias'), ('encoder', 'layer', '2', 'output', 'LayerNorm', 'bias'), ('encoder', 'layer', '1', 'output', 'LayerNorm', 'scale'), ('encoder', 'layer', '4', 'attention', 'output', 'dense', 'bias'), ('encoder', 'layer', '5', 'output', 'LayerNorm', 'bias'), ('encoder', 'layer', '0', 'attention', 'output', 'dense', 'bias'), ('encoder', 'layer', '10', 'attention', 'output', 'dense', 'kernel'), ('encoder', 'layer', '9', 'attention', 'output', 'LayerNorm', 'bias'), ('encoder', 'layer', '8', 'attention', 'output', 'dense', 'kernel'), ('encoder', 'layer', '2', 'attention', 'output', 'dense', 'kernel'), ('encoder', 'layer', '11', 'attention', 'output', 'LayerNorm', 'bias'), ('encoder', 'layer', '4', 'output', 'LayerNorm', 'bias'), ('encoder', 'layer', '3', 'attention', 'output', 'dense', 'kernel'), ('encoder', 'layer', '0', 'output', 'LayerNorm', 'bias'), ('encoder', 'layer', '2', 'attention', 'output', 'LayerNorm', 'scale'), ('encoder', 'layer', '2', 'attention', 'output', 'LayerNorm', 'bias'), ('embeddings', 'LayerNorm', 'scale'), ('encoder', 'layer', '8', 'attention', 'output', 'dense', 'bias'), ('encoder', 'layer', '11', 'output', 'LayerNorm', 'bias'), ('encoder', 'layer', '4', 'attention', 'output', 'dense', 'kernel'), ('encoder', 'layer', '0', 'attention', 'output', 'dense', 'kernel'), ('encoder', 'layer', '11', 'attention', 'output', 'LayerNorm', 'scale'), ('encoder', 'layer', '0', 'output', 'LayerNorm', 'scale'), ('embeddings', 'token_type_embeddings', 'embedding'), ('embeddings', 'LayerNorm', 'bias'), ('encoder', 'layer', '5', 'attention', 'output', 'dense', 'kernel'), ('encoder', 'layer', '11', 'attention', 'output', 'dense', 'kernel'), ('encoder', 'layer', '10', 'output', 'LayerNorm', 'bias'), ('encoder', 'layer', '4', 'attention', 'output', 'LayerNorm', 'bias'), ('encoder', 'layer', '3', 'attention', 'output', 'LayerNorm', 'bias'), ('encoder', 'layer', '7', 'output', 'LayerNorm', 'scale'), ('encoder', 'layer', '0', 'attention', 'output', 'LayerNorm', 'scale'), ('encoder', 'layer', '7', 'attention', 'output', 'dense', 'kernel'), ('encoder', 'layer', '7', 'attention', 'output', 'dense', 'bias'), ('encoder', 'layer', '9', 'attention', 'output', 'dense', 'kernel'), ('encoder', 'layer', '6', 'attention', 'output', 'LayerNorm', 'scale'), ('encoder', 'layer', '1', 'attention', 'output', 'LayerNorm', 'scale'), ('encoder', 'layer', '10', 'attention', 'output', 'LayerNorm', 'bias'), ('encoder', 'layer', '6', 'output', 'LayerNorm', 'scale'), ('encoder', 'layer', '11', 'attention', 'output', 'dense', 'bias'), ('encoder', 'layer', '8', 'output', 'LayerNorm', 'scale'), ('encoder', 'layer', '8', 'output', 'LayerNorm', 'bias'), ('encoder', 'layer', '6', 'attention', 'output', 'dense', 'kernel'), ('encoder', 'layer', '7', 'attention', 'output', 'LayerNorm', 'scale'), ('encoder', 'layer', '7', 'attention', 'output', 'LayerNorm', 'bias'), ('encoder', 'layer', '4', 'output', 'LayerNorm', 'scale'), ('encoder', 'layer', '1', 'attention', 'output', 'dense', 'bias'), ('encoder', 'layer', '1', 'attention', 'output', 'dense', 'kernel'), ('encoder', 'layer', '3', 'attention', 'output', 'dense', 'bias'), ('embeddings', 'word_embeddings', 'embedding'), ('encoder', 'layer', '1', 'attention', 'output', 'LayerNorm', 'bias'), ('encoder', 'layer', '5', 'output', 'LayerNorm', 'scale'), ('encoder', 'layer', '11', 'output', 'LayerNorm', 'scale'), ('encoder', 'layer', '2', 'attention', 'output', 'dense', 'bias'), ('encoder', 'layer', '10', 'attention', 'output', 'dense', 'bias'), ('encoder', 'layer', '3', 'attention', 'output', 'LayerNorm', 'scale'), ('encoder', 'layer', '10', 'attention', 'output', 'LayerNorm', 'scale'), ('encoder', 'layer', '6', 'attention', 'output', 'LayerNorm', 'bias')}
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
```

Apparently almost all parameters are newly initialized, potentially because of different naming conventions e.g. the key `('embeddings', 'word_embeddings', 'weight')` exists but the key `('embeddings', 'word_embeddings', 'embedding')` is expected.

Using `from_pt=True` fixes the issue:

```python
from transformers import FlaxRobertaModel

model = FlaxRobertaModel.from_pretrained("roberta-base", from_pt=True)
```

I don't know if this behavior is intentional, I would have expected an error in that case.

## Expected behavior

I would expect no parameters to be randomly initialized, and the same keys to be unused as when running `RobertaModel.from_pretrained("roberta-base")`.

Here's a colab to reproduce the issue: https://colab.research.google.com/drive/12sDGObc3C1TxW9-k2j4MF69qgpM2kNkP?usp=sharing

Thanks for any help!

Comments:
cc @patrickvonplaten @patil-suraj 

Thank you for reporting this. Yes, the uploaded checkpoints seem to be wrong.

@patrickvonplaten we should verify and re-upload all Roberta flax checkpoints.

Oh oh this doesn't look good at all :-/ 

I'll re-upload the roberta checkpoints. This definitely seems to be a bug

Ok this should be fixed! Thanks a lot for pointing it out @bminixhofer ! I also verified that `roberta-large` and other roberta models work fine.

Works for me too now. Thanks for the quick fix!