Missing `lm_head` parameter in FlaxGPT2LMHeadModel.params
Author: ydshieh
Pull request: no

## Environment info

- `transformers` version: 4.10.0.dev0
- Platform: Linux-5.11.0-25-generic-x86_64-with-glibc2.29
- Python version: 3.8.10
- PyTorch version (GPU?): 1.9.0+cpu (False)
- Tensorflow version (GPU?): 2.5.1 (False)
- Flax version (CPU?/GPU?/TPU?): 0.3.4 (cpu)
- Jax version: 0.2.16
- JaxLib version: 0.1.68
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No

### Who can help
@patrickvonplaten 

## Information

Model I am using: `FlaxGPT2LMHeadModel`

## To reproduce

Steps to reproduce the behavior:

This code snippet 
```
from transformers import FlaxGPT2LMHeadModel

model = FlaxGPT2LMHeadModel.from_pretrained('gpt2')
{k for k in model.params}
```
gives
```
{'transformer'}
```

## Expected behavior

I expect the output will be
```
{'transformer', 'lm_head'}
```

because `FlaxGPT2LMHeadModule` has `self.lm_head` as in the code snippet below

```
class FlaxGPT2LMHeadModule(nn.Module):

    def setup(self):

        self.transformer = FlaxGPT2Module(self.config, dtype=self.dtype)
        self.lm_head = nn.Dense(...)
```

Comments:
Hey @ydshieh,

Thanks for the issue! The reason why `lm_head` is missing in the parameters is because the input and output embeddings are tied for `gpt2` which is the default case if not stated otherwise in the specific config -> see: https://github.com/huggingface/transformers/blob/c71f73f438c7848b7d86af5258e886f03ba45f1e/src/transformers/configuration_utils.py#L227

As a consequence the jax models run through this line of code: https://github.com/huggingface/transformers/blob/c71f73f438c7848b7d86af5258e886f03ba45f1e/src/transformers/models/gpt2/modeling_flax_gpt2.py#L590 which means that the `lm_head` weights are never used **instead** the `shared_kernel = self.transformer.variables["params"]["wte"]["embedding"].T` weights are passed through `lm_head` module. When the first were first created, jax therefore did not trace through an uninitialzied `self.lm_head` but just applied existing weights to the `lm_head` which then didn't create any weights for `lm_head`. 

=> In short one can remember that weights are only created if the control flow goes through the `flax.linen.Module.__call__(...)` method. If the control flow just goes through a `flax.linen.Module.apply(...)` method with given weights, then the model does not expect any weights for this model and will never create it. 

 

Hi, @patrickvonplaten , thank you for this explanation. Now I feel more sure about the code below (copied from `modeling_hybrid_clip.py`) for the recent work on `FlaxEncoderDecoderModel`

```
class FlaxEncoderDecoderModel(FlaxPreTrainedModel):

    @classmethod
    def from_encoder_decoder_pretrained(

        ...

        # init model
        model = cls(config, dtype=dtype, **kwargs)
        model.params["encoder"] = encoder.params
        model.params["decoder"] = decoder.params

        return model
```

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.