Embedding class is replaced when calling `resize_token_embeddings`
Author: TWagner2
Pull request: no

### System Info

- `transformers` version: 4.42.3
- Platform: Linux-5.10.217-205.860.amzn2.x86_64-x86_64-with-glibc2.26
- Python version: 3.11.9
- Huggingface_hub version: 0.23.4
- Safetensors version: 0.4.3
- Accelerate version: not installed
- Accelerate config: not found
- PyTorch version (GPU?): 2.3.1+cu121 (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: No

### Who can help?

@ArthurZucker 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

When calling the `resize_token_embedding` method of MBart, the embeddings of the model are converted from `MBartScaledWordEmbeddings` to just regular `nn.Embeddings`.  This changes the model output because the scaling is lost. Below is a minimal script to demonstrate the issue. In particular, this means that the following series of steps leads to unexpected behavior:
1. Load model, and resize embeddings for a training task
2. Train model
3. Evaluate, results look fine because we both trained and evaluated without scaling in the embeddings
4. Save the model using `save_pretrained` and load it again for production. This will implicitly replace the embedding class with MBartScaledWordEmbeddings again.
5. Run the model. Output is now incorrect because scaling is now included, i.e. the model architecture changed.

I believe other models should also be affected, compare this PR: https://github.com/huggingface/transformers/pull/30410

```python
from transformers.models.mbart import MBartConfig, MBartForCausalLM
import torch

example_input = torch.tensor([1024])

config = MBartConfig(vocab_size=50265, scale_embedding=True)
model = MBartForCausalLM(config)
# This will be an MBartScaledWordEmbedding
print(model.model.decoder.embed_tokens)
old_type = type(model.model.decoder.embed_tokens)
old_output = model.model.decoder.embed_tokens(example_input)
# This will be a standard nn.Embedding
model.resize_token_embeddings(new_num_tokens=50266)
model.resize_token_embeddings(new_num_tokens=50265)
print(model.model.decoder.embed_tokens)
new_type = type(model.model.decoder.embed_tokens)
new_output = model.model.decoder.embed_tokens(example_input)
print(f"Embedding type remained the same: {old_type is new_type}") # False
print(f"Output is the same: {torch.equal(old_output, new_output)}") # False
```

### Expected behavior

The embedding does not change and keeps the scaling.

Comments:
Thanks for reporting the issue! 

For context to whoever might encounter this issue: `MBartScaledWordEmbeddings` was introduced because of backwards compatibility, as we expected users to pass already scaled embeddings in earlier versions (see [comment](https://github.com/huggingface/transformers/pull/30269#discussion_r1570877727)). Imo the prev behavior was a bit misleading but we still needed to keep it.

Solution can be to overwrite `resize_token_embeddings` for these models to account for a custom Embedding module, which will be less breaking. I will add it to my todo list :)

@zucchini-nlp @ArthurZucker 
Can I take this issue up and raise a PR?

@bayllama Sure. Looking forward to your PR. If you have any questions or need assistance, feel free to ask 😄

Thanks a lot for looking at this issue @zucchini-nlp, and thanks @bayllama for taking this up!

> Solution can be to overwrite resize_token_embeddings for these models to account for a custom Embedding module, which will be less breaking.

If I may make a suggestion, in my opinion this would ideally be solved by making the `resize_token__embeddings` method in the base class agnostic to the exact type of embedding, instead of overwriting the method. Otherwise, one runs the risk of re-introducing this bug whenever a new model is added.

@TWagner2 yes, would be great if you can make the `resize_token_embeddings` class agnostic and keep backwards compatibility!

Thanks @zucchini-nlp and @TWagner2 for the inputs. Will try to make the function class agnostic.

Hi @zucchini-nlp and @TWagner2 ,

I took my first stab here in trying to make the function resize_token_embeddings agnostic of the function class and makes sure it always returns the same Custom Embedding Class in the output. I made a commit onto my fork here below with the changes,

`https://github.com/bayllama/transformers/commit/eebabe17909de773aa63e2c57d62974fe7c527be`

Description of the changes:
1) Added **kwargs to the parameters in the init function of the Custom Embedding Class. This allows us the ability to pass the super class parameters to it if needed.
2) Made the padding_idx and Optional[int] and set the default to 0. The reason for this change is while creating an instance of this class in modelling_utils.py we don't want to set a padding_idx.
3) In the modelling_utils.py file, made a change to _get_resized_embeddings where we capture the type of the input class and make sure the output is of the same type.

Please take a look at the commit url above and let me know if I am going in the right direction. Also, if this change is fine I may need to make similar changes to 12 more modeling_{model_name}.py files and make a PR.

Thanks!


@bayllama can you open a PR with the current changes pls? If we're happy with them, then you can modify other models same way

@zucchini-nlp Have made the PR https://github.com/huggingface/transformers/pull/31979

Will need to make changes to the other models that have a custom embedding class also to make sure all tests pass. Let me know

I also experienced this bug, but with different model. With Transformers version 4.43.2, the issue has been solved for `MBartForCausalLM` (like the previous example), but not for `MBartForConditionalGeneration`. Here is the slightly changed code:

```python
from transformers.models.mbart import MBartConfig, MBartForConditionalGeneration
import torch

example_input = torch.tensor([1024])

config = MBartConfig(vocab_size=50265, scale_embedding=True)
model = MBartForConditionalGeneration(config)
# This will be an MBartScaledWordEmbedding
print(model.model.decoder.embed_tokens)
old_type = type(model.model.decoder.embed_tokens)
old_output = model.model.decoder.embed_tokens(example_input)
# This will be a standard nn.Embedding
# model.resize_token_embeddings(new_num_tokens=50266) # This line is optional.
model.resize_token_embeddings(new_num_tokens=50265)
print(model.model.decoder.embed_tokens)
new_type = type(model.model.decoder.embed_tokens)
new_output = model.model.decoder.embed_tokens(example_input)
print(f"Embedding type remained the same: {old_type is new_type}") # False
print(f"Output is the same: {torch.equal(old_output, new_output)}") # False
```

I've tried to analyze the bug. What I found is the method of `MBartModel.get_input_embeddings` returns `self.shared` which has `torch.nn.Embedding` type, shown in [this code](https://github.com/huggingface/transformers/blob/v4.43.2/src/transformers/models/mbart/modeling_mbart.py#L1267):
```python
class MBartModel(MBartPreTrainedModel):
    _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight"]

    def __init__(self, config: MBartConfig):
        super().__init__(config)

        padding_idx, vocab_size = config.pad_token_id, config.vocab_size
        self.shared = nn.Embedding(vocab_size, config.d_model, padding_idx)

        self.encoder = MBartEncoder(config, self.shared)
        self.decoder = MBartDecoder(config, self.shared)

        # Initialize weights and apply final processing
        self.post_init()

    def get_input_embeddings(self):
        return self.shared

    ...
```

However, that type is inconsistent with `MBartEncoder.embed_tokens` which has `MBartScaledWordEmbedding` type, shown in [this code](https://github.com/huggingface/transformers/blob/v4.43.2/src/transformers/models/mbart/modeling_mbart.py#L821): (I haven't analyzed `MBartDecoder` yet.)

```python
class MBartEncoder(MBartPreTrainedModel):
    ...
    def __init__(self, config: MBartConfig, embed_tokens: Optional[nn.Embedding] = None):
        ...
        embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0

        self.embed_tokens = MBartScaledWordEmbedding(
            config.vocab_size, embed_dim, self.padding_idx, embed_scale=embed_scale
        )
        ...
```

And that seems a problem because `MBartModel.get_input_embeddings` only returning the standard embedding type, not the actual embedding which is used by `MBartEncoder`. Furthermore, the method of `resize_token_embeddings` uses `get_input_embeddings` and `set_input_embeddings` as shown in [`PretrainedModel._resize_token_embeddings`](https://github.com/huggingface/transformers/blob/v4.43.2/src/transformers/modeling_utils.py#L1995) implementation:
```python
class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMixin, PeftAdapterMixin):
    ...

    def _resize_token_embeddings(self, new_num_tokens, pad_to_multiple_of=None):
        old_embeddings = self.get_input_embeddings()
        new_embeddings = self._get_resized_embeddings(old_embeddings, new_num_tokens, pad_to_multiple_of)
        ...
        self.set_input_embeddings(new_embeddings)

        ...

    ...
```

But, the persisted embeddings in `MBartEncoder.embed_tokens` is replaced by `self.shared` after resizing because of `MBartModel.set_input_embeddings` method, shown in [this code](https://github.com/huggingface/transformers/blob/v4.43.2/src/transformers/models/mbart/modeling_mbart.py#L1285):
```python
class MBartModel(MBartPreTrainedModel):
    ...

    def set_input_embeddings(self, value):
        self.shared = value
        self.encoder.embed_tokens = self.shared
        self.decoder.embed_tokens = self.shared

    ...
```

I think, to solve this problem, instead of using `torch.nn.Embedding` initialization for `self.shared` in `MBartModel`, let's just use `MBartScaledWordEmbedding`. The implementation is exactly same as the initialization in `MBartEncoder.__init__`.

Can I create another pull request for this solution, or is there any issue for this? 

@AbdiHaryadi I see, the embedding layer in MBartModel has to be the Scaled Embedding also. Yes, feel free to make a PR and tag me when it's ready for review :)

I also found these models need to be aligned with the shared embedding, if you can add it to the PR: BartModel, SeamlessM4TForTextToText, SeamlessM4TForTextToSpeech, SeamlessM4TModel and same for seamless-v2