[Flax] Non working model when exporting to Huggingface
Author: peregilk
Pull request: no

I have trained a RoBERTa base Norwegian according to instructions given at https://github.com/huggingface/transformers/tree/master/examples/flax/language-modeling#masked-language-modeling.

The final mlm accuracy is 0.63, indicating a working model.

I am trying to load the model, and export til to PyTorch (or TF) for using the inference widget on Hugging Face.

The following code runs without errors:

from transformers import AutoTokenizer, RobertaForMaskedLM
model = RobertaForMaskedLM.from_pretrained('model_dir', from_flax=True)
tokenizer = AutoTokenizer.from_pretrained('model_dir')
model.save_frompretrained('.')
tokenizer.save_frompretrained('.')
Example widget here: https://huggingface.co/pere/norwegian-roberta-base?text=Dette+er+en+%3Cmask%3E.

The outputs makes absolutely no sense. What is the correct way of exporting a Flax model (with and without the MLM head)?

Comments:
cc @patrickvonplaten @patil-suraj 

Hi @peregilk, how did you save the flax Roberta model using the `FlaxRobertaModel` class or using the `FlaxRobertaForMaskedLM` class? 
I tried loading the flax model with
```python
fx_model = FlaxRobertaForMaskedLM.from_pretrained("pere/norwegian-roberta-base")
```
and it gives this warning 
```
Some weights of FlaxRobertaForMaskedLM were not initialized from the model checkpoint at pere/norwegian-roberta-base and are newly initialized: {('lm_head', 'dense', 'kernel'), ('lm_head', 'layer_norm', 'bias'), ('lm_head', 'layer_norm', 'scale'), ('lm_head', 'bias'), ('lm_head', 'dense', 'bias')}
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
```

which means the model was not saved with `FlaxRobertaForMaskedLM`, so the `lm_head` is missing. In which case it is randomly initialized, which might explain the above behavior.

Please make sure you use the `*ForMaskedLM` class to save, load MLM models. Thanks!

Thanks for the feedback @patil-suraj. I think you are absolutely right that the uploaded model does not contain the LM head. The issue is then probably related to the loading.

I have tried both these alternatives:
```
model = RobertaForMaskedLM.from_pretrained('model_dir', from_flax=True)
model = FlaxRobertaForMaskedLM.from_pretrained('model_dir')
```

Both giving roughly the same "random" result. I actually notice errors from both these methods when loading. Slightly different, but both pointing at the lm head not being loaded correctly:

```
- This IS NOT expected if you are initializing RobertaForMaskedLM from a Flax model that you expect to be exactly identical (e.g. initializing a BertForSequenceClassification model from a FlaxBertForSequenceClassification model).
This is the version that is currently pushed, and you are probably right the head is not included.
```


```
Some weights of FlaxRobertaForMaskedLM were not initialized from the model checkpoint at . and are newly initialized: {('lm_head', 'layer_norm', 'scale'), ('lm_head', 'bias'), ('lm_head', 'dense', 'kernel'), ('lm_head', 'layer_norm', 'bias'), ('lm_head', 'dense', 'bias')}
```

How can I load the Flax model with the LM head?


Related: How do I check that `flax_model.msgpack` actually is saved with the LM head?

@peregilk - I seems like the uploaded flax model weights are not correct. When running:

```python
from transformers import FlaxRobertaForMaskedLM
model = FlaxRobertaForMaskedLM.from_pretrained("pere/norwegian-roberta-base")
```

The lm weights are missing. 

I think this commit: https://huggingface.co/pere/norwegian-roberta-base/commit/163f3993e02f514ce29901020a7ec73958fb415b was where they were deleted. Can you revert this change and instead use the Flax model weights from this commit: https://huggingface.co/pere/norwegian-roberta-base/commit/506fb71065742c28f1efe8ccb9cae48ddcc562ad .

If you convert the flax weights of this commit: https://huggingface.co/pere/norwegian-roberta-base/commit/506fb71065742c28f1efe8ccb9cae48ddcc562ad to a `RobertaForMaskedLM` pytorch model, the widget should work correctly

@patrickvonplaten. Awesome! Works perfectly! I must have overwritten this in my first attempts. Did not realise that loading and saving FlaxRobertaModel removes the LM head. But it makes perfectly sense that it does.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.