[FlaxBartForCausalLM] Embed tokens not loaded in Flax decoder model from encoder-decoder weights
Author: sanchit-gandhi
Pull request: no

The embeddings module `embed_tokens` is not loaded from pre-trained Flax model weights when a FlaxBartForCausalLM model is instantiated in isolation. As a consequence, these embedding weights are randomly initialised. The following code snippet demonstrates this fact by comparing the FlaxBartForCausalLM model to its PyTorch equivalent, BartForCausalLM. For the PyTorch (resp. Flax) model, the weights are loaded from pre-trained PyTorch (resp. Flax) weights at https://huggingface.co/sanchit-gandhi/tiny-random-bart. These model weights are identical to those in the repository at https://huggingface.co/hf-internal-testing/tiny-random-bart, but with the exception that this repository contains both Flax and PyTorch weights, unlike those at hf-internal-testing which contain only PyTorch weights.

```python
from transformers import BartForCausalLM, FlaxBartForCausalLM
import tempfile
from flax.traverse_util import flatten_dict

pt_dec_model = BartForCausalLM.from_pretrained('sanchit-gandhi/tiny-random-bart')
fx_dec_model = FlaxBartForCausalLM.from_pretrained('sanchit-gandhi/tiny-random-bart')

# Convert the PT model to FX 
with tempfile.TemporaryDirectory() as tmpdirname:
    pt_dec_model.save_pretrained(tmpdirname)
    pt_dec_model_to_fx = FlaxBartForCausalLM.from_pretrained(tmpdirname, from_pt=True)
    
# easier to work in terms of flattened dicts
pt_dec_params_to_fx = flatten_dict(pt_dec_model_to_fx.params)
fx_dec_params = flatten_dict(fx_dec_model.params)

# Check that all keys match
assert fx_dec_params.keys() == pt_dec_params_to_fx.keys()

# Check that all the weights are **precisely** equal
for param in pt_dec_params_to_fx:
    assert (fx_dec_params[param] == pt_dec_params_to_fx[param]).all(), param
```

Output:
```
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Input In [211], in <cell line: 21>()
     20 # Check that all the weights are **precisely** equal
     21 for param in pt_dec_params_to_fx:
---> 22     assert (fx_dec_params[param] == pt_dec_params_to_fx[param]).all(), param

AssertionError: ('model', 'decoder', 'embed_tokens', 'embedding')
```
We see here that the embedding weights do not match for the standalone decoder models: the `embed_tokens` are not loaded from pre-trained in Flax, and are instead randomly initialised.

Loading full encoder-decoder models, we see that the weights match for the embeddings:

```python
from transformers import BartModel, FlaxBartModel
import tempfile
from flax.traverse_util import flatten_dict

pt_model = BartModel.from_pretrained('sanchit-gandhi/tiny-random-bart')
fx_model = FlaxBartModel.from_pretrained('sanchit-gandhi/tiny-random-bart')

# Convert the PT model to FX 
with tempfile.TemporaryDirectory() as tmpdirname:
    pt_model.save_pretrained(tmpdirname)
    pt_model_to_fx = FlaxBartModel.from_pretrained(tmpdirname, from_pt=True)
    
# easier to work in terms of flattened dicts
pt_params_to_fx = flatten_dict(pt_model_to_fx.params)
fx_params = flatten_dict(fx_model.params)

# Check that all keys match
assert fx_params.keys() == pt_params_to_fx.keys()

# Check that all the weights are **precisely** equal
for param in pt_params_to_fx:
    assert (fx_params[param] == pt_params_to_fx[param]).all(), param
```
A fix is needed to be able to load the Flax encoder-decoder embedding weights into a standalone decoder module.

Comments:
In the PyTorch Bart modelling script, we first define a 'shared' nn.Embedding module, which we then directly pass to the encoder and decoder modules to explicitly tie their word embeddings:
https://github.com/huggingface/transformers/blob/14daa6102a0e8a35ef734dd21bfcf31d9b0207d1/src/transformers/models/bart/modeling_bart.py#L1146-L1149
Due to the stateful nature of PyTorch modules, we can then overwrite this embedding in the `init` method of the encoder or decoder, depending on whether or not the optional keyword argument `embed_tokens` is specified:
https://github.com/huggingface/transformers/blob/cc034f72eb6137f4c550e911fba67f8a0e1e98fa/src/transformers/models/bart/modeling_bart.py#L710-L713
(Note that for decoder-only models, we do not specify the argument `embed_tokens` for the decoder module. Thus, it defaults to being initialised in the decoder module's `init`). For the encoder-decoder model, there are three instances in which the embeddings are defined: as `shared` under the BartModel, and again as `embed_tokens` in the encoder and decoder models. This yields the following parameter tree:
```
PT enc-dec model
    shared
    encoder
        embed_tokens
        ...
    decoder
        embed_tokens
        ...
```
Likewise, in the Flax Bart modelling script, we first define a 'shared' nn.Embed module, which we then directly pass to the encoder and decoder modules to explicitly tie their word embeddings:
https://github.com/huggingface/transformers/blob/cc034f72eb6137f4c550e911fba67f8a0e1e98fa/src/transformers/models/bart/modeling_flax_bart.py#L839-L847
However, due to the stateless nature of JAX/Flax models, we cannot then overwrite this embedding in the `setup` method of the encoder or decoder. To address this, it was decided in #15920 that the keyword argument `embed_tokens` must always be specified to the encoder/decoder modules. Thus, there is only one instance in which the embeddings are defined: as `shared` under the FlaxBartModel. This results in different parameter tree to that in PyTorch:
```
FX enc-dec model
    shared
    encoder
        ...
    decoder
        ...
```
For encoder-decoder models, PyTorch to Flax conversion is possible: the Flax encoder-decoder model is able to leverage the PyTorch `shared` embedding weights, and then pass these into the encoder and decoder separately (effectively tying the weights, but only having one variable) .

However, an issue arises for decoder only models. Here, the Flax decoder cannot leverage all of the Flax encoder-decoder model weights. This is due to the format of its parameter tree, which is constructed jointly through the [FlaxBartDecoderWrapper](https://github.com/huggingface/transformers/blob/a192f61e0825150e54e15fdc451cf37e23532b3f/src/transformers/models/bart/modeling_flax_bart.py#L1863) and [FlaxBartForCausalLM](https://github.com/huggingface/transformers/blob/a192f61e0825150e54e15fdc451cf37e23532b3f/src/transformers/models/bart/modeling_flax_bart.py#L1885) module:
```
FX dec-only model
    decoder
        embed_tokens
        ...
```
 Since there is no module `shared` in Flax decoder only models, the system is not able to leverage the embedding weights registered under `shared` from the Flax encoder-decoder model weights. However, `embed_tokens` is now defined under the decoder module, meaning that we are able to leverage PyTorch encoder-decoder or decoder-only model weights and load them into Flax:
```
PT dec-only model
    decoder
        embed_tokens
        ...
```

Potential solutions:
- In the [FlaxBartDecoderWrapper](https://github.com/huggingface/transformers/blob/a192f61e0825150e54e15fdc451cf37e23532b3f/src/transformers/models/bart/modeling_flax_bart.py#L1863), we can rename `embed_tokens` to `self.shared`, thus bringing the param trees of the Flax encoder-decoder and decoder-only models into alignment. Doing so enables the decoder only embeddings to be loaded from Flax encoder-decoder model weights. However, this is not an ideal solution: by renaming the module, we will no longer be able to load Flax decoder only model weights from PyTorch (encoder-)decoder weights, as these parameter trees will not match.
- We could define a `self.shared` nn.Embedding module in the PyTorch [DecoderWrapper](https://github.com/huggingface/transformers/blob/a192f61e0825150e54e15fdc451cf37e23532b3f/src/transformers/models/bart/modeling_bart.py#L1680) and then pass this into the decoder model. This maintains consistency between the encoder-decoder style models and the decoder-only ones. (Define `shared` outside the modules, then pass it in, thus registering a state-dict of  `(shared), (decoder, embed_tokens)` instead of just `(decoder, embed_tokens)`). However, this is a breaking change for PyTorch Bart models, and should be avoided.
- What is probably easier and more effective than both of the above is explicit naming of the `embed_tokens` module in the Flax Bart encoder and decoder modules, giving a parameter tree that exactly matches the PyTorch one, both for encoder-decoder and decoder only models.
```
Enc-dec model
    shared
    encoder
        embed_tokens
        ...
    decoder
        embed_tokens
        ...
```
```
Dec-only model
    decoder
        embed_tokens
        ...
```
It is not apparent from the Flax docs how this explicitly naming can be achieved, but I have asked the Flax community on the Flax discussions page how to go about doing this: https://github.com/google/flax/discussions/2046#discussion-4004536

cc @patrickvonplaten @patil-suraj 

As a temporary fix, a standalone Flax decoder model can be loaded entirely from it's equivalent PyTorch weights and the `embed_tokens` made to match:

```python
from transformers import BartForCausalLM, FlaxBartForCausalLM
import tempfile
from flax.traverse_util import flatten_dict

pt_dec_model = BartForCausalLM.from_pretrained('sanchit-gandhi/tiny-random-bart')
# force Flax weights to be loaded from PyTorch - enables `embed_tokens` to be loaded correctly
fx_dec_model = FlaxBartForCausalLM.from_pretrained('sanchit-gandhi/tiny-random-bart', from_pt=True)

# Convert the PT model to FX 
with tempfile.TemporaryDirectory() as tmpdirname:
    pt_dec_model.save_pretrained(tmpdirname)
    pt_dec_model_to_fx = FlaxBartForCausalLM.from_pretrained(tmpdirname, from_pt=True)
    
# easier to work in terms of flattened dicts
pt_dec_params_to_fx = flatten_dict(pt_dec_model_to_fx.params)
fx_dec_params = flatten_dict(fx_dec_model.params)

# Check that all keys match
assert fx_dec_params.keys() == pt_dec_params_to_fx.keys()

# Check that all the weights are **precisely** equal
for param in pt_dec_params_to_fx:
    assert (fx_dec_params[param] == pt_dec_params_to_fx[param]).all(), param
```

Hmm isn't the problem here that the weights are not correctly mapped? E.g. when I run the first part of your codesnippet:

```py
from transformers import BartForCausalLM, FlaxBartForCausalLM
import tempfile
from flax.traverse_util import flatten_dict

pt_dec_model = BartForCausalLM.from_pretrained('sanchit-gandhi/tiny-random-bart')
fx_dec_model = FlaxBartForCausalLM.from_pretrained('sanchit-gandhi/tiny-random-bart')

# Convert the PT model to FX 
with tempfile.TemporaryDirectory() as tmpdirname:
    pt_dec_model.save_pretrained(tmpdirname)
    pt_dec_model_to_fx = FlaxBartForCausalLM.from_pretrained(tmpdirname, from_pt=True)
```

It says that:

```
...
Some weights of FlaxBartForCausalLM were not initialized from the model checkpoint at sanchit-gandhi/tiny-random-bart and are newly initialized: {('model', 'decoder', 'embed_tokens', 'embedding')}
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
Some weights of the model checkpoint at /tmp/tmpqaqqrtft were not used when initializing FlaxBartForCausalLM: {('lm_head', 'kernel')}
- This IS expected if you are initializing FlaxBartForCausalLM 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 FlaxBartForCausalLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
```

which shows that those weights are not correctly loaded into the Flax model, so IMO the bug is that the Flax model doesn't correctly convert the naming here. 

Also note that 99% of all Bart models have their output embeddings tied to their input embeddings, see: https://github.com/huggingface/transformers/blob/eb5bdcdfa51f743887ee1d9c7f230444d7a8b23c/src/transformers/models/bart/modeling_flax_bart.py#L1929 in which case the `lm_head` weights are irrelevant. But agree that there is an error nevertheless.

The solution should be to correct the naming / weight conversion here though IMO

The initialisation of parameters works slightly differently between PyTorch and Flax. In PyTorch, any module defined under a model's `init` will be added to the model's state-dict. In Flax, modules are first defined in the `setup` method, but are only added to the param dict if traced in the `call` method when the dummy forward pass is performed during model initialisation. 

With that being said, the `lm_head` is always added to the state-dict in PyTorch, whether or not the word embeddings are tied. However, this is not the case in Flax - the `lm_head` is only added to the param dict _if_ used in the `call` method. Inspecting the model code, we see this is only the case if the word embeddings are not tied:
https://github.com/huggingface/transformers/blob/eb5bdcdfa51f743887ee1d9c7f230444d7a8b23c/src/transformers/models/bart/modeling_flax_bart.py#L1927-L1931

If we look back to the code snippet for the loading of the decoder-only PyTorch and Flax models, we can confirm that the word embeddings are tied, and that the `lm_head` is instantiated for the PyTorch model and not the Flax one (as expected):
```python
from transformers import BartForCausalLM, FlaxBartForCausalLM

pt_dec_model = BartForCausalLM.from_pretrained('sanchit-gandhi/tiny-random-bart')
fx_dec_model = FlaxBartForCausalLM.from_pretrained('sanchit-gandhi/tiny-random-bart')

print(f"Tie word embeddings? PyTorch: {pt_dec_model.config.tie_word_embeddings}, Flax: {fx_dec_model.config.tie_word_embeddings}")
print(f"PyTorch Decoder modules: {[n for n, _ in pt_dec_model.named_children()]}")
print(f"Flax Decoder modules: {fx_dec_model.params.keys()}")
```
Output:
```
...
Some weights of FlaxBartForCausalLM were not initialized from the model checkpoint at sanchit-gandhi/tiny-random-bart 
and are newly initialized: {('model', 'decoder', 'embed_tokens', 'embedding')}
...
Tie word embeddings? PyTorch: True, Flax: True
PyTorch Decoder modules: ['model', 'lm_head']
Flax Decoder modules: dict_keys(['model'])
```
When loading from pre-trained Flax weights, we see that the only parameters randomly initialised are the `embed_tokens`.

We perform a slightly different operation when comparing the PyTorch weights to those in Flax - we first save the PyTorch model to a temporary directory (`.save_pretrained(tmpdirname)`) and then load this model from it's PyTorch weights into Flax: 
```python
# Convert the PT model to FX 
with tempfile.TemporaryDirectory() as tmpdirname:
    pt_dec_model.save_pretrained(tmpdirname)
    pt_dec_model_to_fx = FlaxBartForCausalLM.from_pretrained(tmpdirname, from_pt=True)
```
Here, the PyTorch weights for the `lm_head` are saved. However, since the `lm_head` is not used in the `call` method of the Flax model,  they are subsequently not used when loading the PyTorch model into Flax. Thus, we expect to see the aforementioned message:
```
Some weights of the model checkpoint at /tmp/tmpqaqqrtft were not used when initializing FlaxBartForCausalLM: {('lm_head', 'kernel')}
```

When we run the full code snippet, we see that the only weights that do not match between PyTorch and Flax decoder-only models are the `embed_tokens` - the `lm_head` is not used in Flax and so is ignored from this comparison. This is the core issue, which arises due to a different parameter structure between the Flax encoder-decoder models and the Flax decoder-only models.

For Flax encoder-decoder, the tied word embeddings are held under the module `shared`, which explicitly ties the word embedding tokens for the encoder and decoder:
```
FX enc-dec model
    shared
    encoder
        ...
    decoder
        ...
```
For Flax decoder-only models, we do not have the module `shared`, giving the modified parameter tree:
```
FX dec-only model
    decoder
        embed_tokens
        ...
```
The reason we omit the `shared` module is to give one-to-one equivalence to the corresponding PyTorch decoder-only state-dict:
```
PT dec-only model
    decoder
        embed_tokens
        ...
```
To remedy this issue, we have three choices:
1. We can either insert a module named `shared` for the Flax decoder-only, and enable it to be compatible with Flax encoder-decoder models:
```
FX dec-only model
    shared
    decoder
        ...
```
However, this would then break equivalence between PT and FX decoder-only models, the parameter trees now differing.
2. We keep the current structure and allow for PT and FX decoder-only model equivalence.
3. We explicitly add `embed_tokens` as a named module under the FX encoder-decoder model:
```
FX enc-dec model
    shared
    encoder
        embed_tokens
        ...
    decoder
        embed_tokens
        ...
```
Which enables PT - FX equivalence for both encoder-decoder and decoder-only models.

Of the three, the latter is my preference, as it allows for full compatibility between the different frameworks.

The full code snippet that examines the `tie_word_embeddings` variable as well as the parameter weights:
```python
from transformers import BartForCausalLM, FlaxBartForCausalLM

pt_dec_model = BartForCausalLM.from_pretrained('sanchit-gandhi/tiny-random-bart')
fx_dec_model = FlaxBartForCausalLM.from_pretrained('sanchit-gandhi/tiny-random-bart')

print(f"Tie word embeddings? PyTorch: {pt_dec_model.config.tie_word_embeddings}, Flax: {fx_dec_model.config.tie_word_embeddings}")
print(f"PyTorch Decoder modules: {[n for n, _ in pt_dec_model.named_children()]}")
print(f"Flax Decoder modules: {fx_dec_model.params.keys()}")

# Convert the PT model to FX 
with tempfile.TemporaryDirectory() as tmpdirname:
    pt_dec_model.save_pretrained(tmpdirname)
    pt_dec_model_to_fx = FlaxBartForCausalLM.from_pretrained(tmpdirname, from_pt=True)
    
# easier to work in terms of flattened dicts
pt_dec_params_to_fx = flatten_dict(pt_dec_model_to_fx.params)
fx_dec_params = flatten_dict(fx_dec_model.params)

# Check that all keys match
assert fx_dec_params.keys() == pt_dec_params_to_fx.keys()

# Check that all the weights are **precisely** equal
mismatch_params = []
print("Checking weights match...")
for param in pt_dec_params_to_fx:
    if (fx_dec_params[param] != pt_dec_params_to_fx[param]).all():
        mismatch_params.append(param)
if len(mismatch_params) == 0:
    print("✅ All PyTorch and Flax parameters match")
else:
    print("❌ The following weights do not match:")
    for param in mismatch_params:
        print(param)
```
Output:
```
Tie word embeddings? PyTorch: True, Flax: True
PyTorch Decoder modules: ['model', 'lm_head']
Flax Decoder modules: dict_keys(['model'])
Checking weights match...
❌ The following weights do not match:
('model', 'decoder', 'embed_tokens', 'embedding')
```

cc @patil-suraj here since he added `bart-large`

I'm fine with whatever solution as this is really an edge case - however we should not break backward compatibility here especially with respect to the weights structure. Also we should **not** touch the PyTorch Bart code

Agree that we should not change the PyTorch modelling code! My preference is modifying the Flax encoder-decoder param dict to explicitly include `embed_tokens` under the encoder and decoder modules (as with the PyTorch models and the Flax decoder-only models) which will bring compatibility between all four models (PyTorch encoder-decoder, Flax encoder-decoder, PyTorch decoder-only and Flax decoder-only)

Ok for me! @patil-suraj what do you think?

I'm fine with modifying the param dict of `bart`  here, since flax doesn't add those `embed_tokens` weights under `encoder` and `decoder` if they are initialised outside and shared. 

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/main/CONTRIBUTING.md) are likely to be ignored.

From the discussion with the Flax authors at https://github.com/google/flax/discussions/2046#discussion-4004536, the best option appears to be handling this in Flax weights loading script.

Generally, not very keen on changing the general Flax weight conversion script because of only a single model. But happy to iterate over the design in a PR. @sanchit-gandhi, could you maybe open a PR to show how you would like to solve the problem and then we take it from there? 

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/main/CONTRIBUTING.md) are likely to be ignored.