Fix Embedding Module Bug in Flax Models
Author: sanchit-gandhi
Pull request: yes

This PR fixes a widespread bug relating to the way in which Embedding Modules are defined for a number of different Flax Modules. In all of these instances, `embed_tokens` is defined as an optional `nn.Embed` attribute in the corresponding Flax Module:
https://github.com/huggingface/transformers/blob/b693cbf99c5a180dde8b32ded2fb82ea735aab15/src/transformers/models/bart/modeling_flax_bart.py#L698-L701
When `embed_tokens` is specified as an `nn.Embed` module and passed as an argument to said Flax module, there are no issues. However, when `embed_tokens` is omitted from the arguments, it defaults to `None`. The `call` method in the Flax Module then attempts to override this attribute with an `nn.Embed` module:
https://github.com/huggingface/transformers/blob/b693cbf99c5a180dde8b32ded2fb82ea735aab15/src/transformers/models/bart/modeling_flax_bart.py#L711-L716
Modifying a submodule's attributes after constructing it violates the _stateless_ design philosophy adopted by Flax. Doing so results in a [`SetAttributeFrozenModuleError`](https://flax.readthedocs.io/en/latest/flax.errors.html#flax.errors.SetAttributeFrozenModuleError). This PR addresses this issue by forcing `embed_tokens` to be a required argument for the Flax Module, thus avoiding the need to ever override the attribute.

There is a requirement to keep the Embedding Module defined as `embed_tokens` in order to remain consistent with the analogous PyTorch models and facilitate the conversion of Flax models to PyTorch.

For reference, the alternative workarounds are not feasible:

1. Defining a new `nn.Embed` module `embeddings` in `setup` and setting `name=embed_tokens`

2. Defining a new `nn.Embed` module `embeddings` in `setup` and passing the `embed_tokens` as parameters in `module.apply`

For 1, it is not possible to define two modules with the same name. For 2, the module name differs from the PyTorch script, and so Flax to PyTorch conversion is not possible.

Currently, the modified Flax Modules are only used in an Encoder-Decoder configuration, where the embeddings of the encoder are tied to the decoder. Hence, in all instances, `embed_tokens` is specified as an `nn.Embed` argument _within_ the Module definition, and thus cannot be partioned. Since there are no standalone Encoder/Decoder models, this edge-case cannot yet be tested. Developing standalone Causal LMs in Flax is currently a WIP. Once complete, there will be a suite of models with which to run these edge-case tests.

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_15920). All of your documentation changes will be reflected on that endpoint.

@sanchit-gandhi thanks for the PR - could you add a code snippet that shows how this design leads to an error? Not a big fan in general of adding `nn.Module` as class attributes of another `nn.Module` - so not sure if this PR is the correct way. @patil-suraj what do you think?

Of course @patrickvonplaten! Here's an excerpt of code that demonstrates the [`SetAttributeFrozenModuleError`](https://flax.readthedocs.io/en/latest/flax.errors.html#flax.errors.SetAttributeFrozenModuleError) described above. Here, `embed_tokens` is passed as an argument to the `encoder`, thus enabling it to be initialised correctly. Conversely, `embed_tokens` is omitted as an argument to the `decoder`, and so defaults to `None`. When the `decoder` is initialised, Flax attempts to override the model attribute `embed_tokens` that has already been constructed as a default argument (`None`). This causes the script to fail, as the module attributes are frozen after construction:
```python
import flax.linen as nn
from transformers import BartConfig
import jax
import jax.numpy as jnp
from transformers.models.bart.modeling_flax_bart import FlaxBartEncoder, FlaxBartDecoder, FlaxBartPreTrainedModel

class FlaxDummyBartModule(nn.Module):
    config: BartConfig
    dtype: jnp.dtype = jnp.float32  # the dtype of the computation

    def setup(self):
        self.embedding = nn.Embed(
            self.config.vocab_size,
            self.config.d_model,
            embedding_init=jax.nn.initializers.normal(self.config.init_std),
        )

        self.encoder = FlaxBartEncoder(self.config, dtype=self.dtype, embed_tokens=self.embedding)
        self.decoder = FlaxBartDecoder(self.config, dtype=self.dtype)

    def __call__(
        self,
        input_ids,
        attention_mask,
        decoder_input_ids,
        decoder_attention_mask,
        position_ids,
        decoder_position_ids,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = True,
        deterministic: bool = True,
    ):
        encoder_outputs = self.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            deterministic=deterministic,
        )

        decoder_outputs = self.decoder(
            input_ids=decoder_input_ids,
            attention_mask=decoder_attention_mask,
            position_ids=decoder_position_ids,
            encoder_hidden_states=encoder_outputs[0],
            encoder_attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            deterministic=deterministic,
        )

        return decoder_outputs

class FlaxDummyBartModel(FlaxBartPreTrainedModel):
    config: BartConfig
    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
    module_class = FlaxDummyBartModule

model = FlaxDummyBartModel.from_pretrained('hf-internal-testing/tiny-random-bart', from_pt=True)
```
<details open>
<summary> Output </summary>
<br>

```python
SetAttributeInModuleSetupError            Traceback (most recent call last)
Input In [4], in <cell line: 1>()
----> 1 FlaxDummyBartModel.from_pretrained('hf-internal-testing/tiny-random-bart', from_pt=True)

File ~/transformers/src/transformers/modeling_flax_utils.py:550, in FlaxPreTrainedModel.from_pretrained(cls, pretrained_model_name_or_path, dtype, *model_args, **kwargs)
    547     resolved_archive_file = None
    549 # init random models
--> 550 model = cls(config, *model_args, **model_kwargs)
    552 if from_pt:
    553     state = load_pytorch_checkpoint_in_flax_state_dict(model, resolved_archive_file)

File ~/transformers/src/transformers/models/bart/modeling_flax_bart.py:933, in FlaxBartPreTrainedModel.__init__(self, config, input_shape, seed, dtype, **kwargs)
    924 def __init__(
    925     self,
    926     config: BartConfig,
   (...)
    930     **kwargs
    931 ):
    932     module = self.module_class(config=config, dtype=dtype, **kwargs)
--> 933     super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype)

File ~/transformers/src/transformers/modeling_flax_utils.py:116, in FlaxPreTrainedModel.__init__(self, config, module, input_shape, seed, dtype)
    113 self.dtype = dtype
    115 # randomly initialized parameters
--> 116 random_params = self.init_weights(self.key, input_shape)
    118 # save required_params as set
    119 self._required_params = set(flatten_dict(unfreeze(random_params)).keys())

File ~/transformers/src/transformers/models/bart/modeling_flax_bart.py:963, in FlaxBartPreTrainedModel.init_weights(self, rng, input_shape)
    953     module_init_outputs = self.module.init(
    954         rngs,
    955         input_ids,
   (...)
    960         return_dict=False,
    961     )
    962 else:
--> 963     module_init_outputs = self.module.init(
    964     rngs,
    965     input_ids,
    966     attention_mask,
    967     decoder_input_ids,
    968     decoder_attention_mask,
    969     position_ids,
    970     decoder_position_ids,
    971 )
    972 return module_init_outputs["params"]

    [... skipping hidden 11 frame]

Input In [2], in FlaxDummyBartModule.__call__(self, input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, position_ids, decoder_position_ids, output_attentions, output_hidden_states, return_dict, deterministic)
     15 def __call__(
     16     self,
     17     input_ids,
   (...)
     26     deterministic: bool = True,
     27 ):
     28     encoder_outputs = self.encoder(
     29         input_ids=input_ids,
     30         attention_mask=attention_mask,
   (...)
     35         deterministic=deterministic,
     36     )
---> 38     decoder_outputs = self.decoder(
     39         input_ids=decoder_input_ids,
     40         attention_mask=decoder_attention_mask,
     41         position_ids=decoder_position_ids,
     42         encoder_hidden_states=encoder_outputs[0],
     43         encoder_attention_mask=attention_mask,
     44         output_attentions=output_attentions,
     45         output_hidden_states=output_hidden_states,
     46         return_dict=return_dict,
     47         deterministic=deterministic,
     48     )
     50     return decoder_outputs

    [... skipping hidden 6 frame]

File ~/transformers/src/transformers/models/bart/modeling_flax_bart.py:783, in FlaxBartDecoder.setup(self)
    780 self.embed_scale = math.sqrt(self.config.d_model) if self.config.scale_embedding else 1.0
    782 if self.embed_tokens is None:
--> 783     self.embed_tokens = nn.Embed(
    784         self.config.vocab_size,
    785         embed_dim,
    786         embedding_init=jax.nn.initializers.normal(self.config.init_std),
    787     )
    789 # Bart is set up so that if padding_idx is specified then offset the embedding ids by 2
    790 # and adjust num_embeddings appropriately. Other models don't have this hack
    791 self.offset = 2

File ~/venv/lib/python3.8/site-packages/flax/linen/module.py:673, in Module.__setattr__(self, name, val)
    671 if is_dataclass_attr:
    672   if self._state.in_setup:
--> 673     raise errors.SetAttributeInModuleSetupError()
    674   object.__setattr__(self, name, val)
    675 # Submodules are being defined and attached in setup()
    676 else:

SetAttributeInModuleSetupError: Module construction attributes are frozen. (https://flax.readthedocs.io/en/latest/flax.errors.html#flax.errors.SetAttributeInModuleSetupError)
```
</details>

With regards to adding an `nn.Module` as a class attribute to another `nn.Module`, there are several pre-existing models that employ this design philosophy, including the FlaxBartModule itself:
https://github.com/huggingface/transformers/blob/9932ee4b4bca9045d941af6687ef69eedcf68483/src/transformers/models/bart/modeling_flax_bart.py#L855-L862
I too agree that it is not the most elegant solution. However, defining `embed_tokens` as an `nn.Module` and passing it as an argument to another `nn.Module` does ensure consistency between the parameter naming in Flax and PyTorch. By extension, it facilitates model conversion between Flax and PyTorch, which is crucial.


I agree with this approach. Had an offline discussion about this with @sanchit-gandhi.
Since flax module attributes are frozen, it's impossible to initialise and set embeddings if the embedding is defined as a module attribute. This is currently a bug in all FlaxBart like models. As demonstrated by the above example.

And to keep naming consistency with PT, we cannot introduce a new name for these embeddings. I am not sure if there's any other way to avoid this. So I am okay with this PR.




Thanks for the code snippet @sanchit-gandhi. If you're ok with it @patil-suraj - it's good to go for me, I trust your judgement here! 

All encoder-decoder models pass the embedding weights from the higher level class to the encoder decoder, so I assume that there is no pretrained model that has encoder or decoder embedding weights directly saved in the dictionary, which is why this PR should have 0 breaking changes, right @patil-suraj ? 



> this PR should have 0 breaking changes, right @patil-suraj ?

@patrickvonplaten yes, you are right. Also just ran slow tests for BART to confirm this and they all pass.