Modify resize_token_embeddings to ensure output type is same as input
Author: bayllama
Pull request: yes

# What does this PR do?

Modified resize_token_embeddings to return the same class that is passed as input to it. Today, even if a custom embedding class is passed, resize_token_embeddings converts it to a nn.Embedding but this commit makes sure that does not happen and the custom embedding class is returned.
<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes # (31835)
[](https://github.com/huggingface/transformers/issues/31835)

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?
@zucchini-nlp


<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @


 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker
- vision models: @amyeroberts
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @SunMarc

Integrations:

- deepspeed: HF Trainer/Accelerate: @muellerzr
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc 

Documentation: @stevhliu and @MKhalusova

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
@zucchini-nlp 
That's a good idea. I have made the change you recommended.

@zucchini-nlp @amyeroberts Added the comment describing the change

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_31979). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@bayllama the tests are failing because we pass `device` and `dtype` to the Embedding Class when resizing, but the custom classes don't accept any kwargs. Can you add `**kwargs` plz to init in the correct device and dtype?

You should be able to verify all tests are passing by this command :)
```bash
pytest -k test_resize_token tests/models/
```

@zucchini-nlp 

When I run the pytest that you suggested I don't see the error that you have mentioned above. We don't set the device and dtype for the Custom Embedding Class at all right?

```
new_embeddings = nn.Embedding(
            new_num_tokens,
            old_embedding_dim,
            device=old_embeddings.weight.device,
            dtype=old_embeddings.weight.dtype,
        )
```

We only do it for the nn.Embedding and the replace the weight.data in the old embeddings which is already in the correct device and has the right dtype.

However I do see another issue which is in hte resize_token_embeddings, this attribute "text_config" is being modified and I haven't accounted for that in my previous commit. I will work on a fix for this. However can you please give me more details on the dtype and device issue, because I don't see it in my set of tests.



@bayllama my bad, I was running in the a different branch. You're right, some VLMs are failing, as I see it should be fixed with one line after we swapped weights! Let me know when it's fixed :) 

@zucchini-nlp Changing the shape of the old_embeddings would fix this, however shape attribute is not writable in torch, hence we cannot do something like this,
```
old_embeddings.weight.shape = new_embeddings.weight.shape
```

Hence I am thinking of what would be the best way to do this. If nothing works out I may need to go to the method I was following in my first PR where I check the type of input embedding and create a new object and return it.

@bayllama I see, but we can also change the `num_embeddings` attribute which is writable.
Like `old_embeddings.num_embeddings = new_embeddings.weight.data.shape[0]`

@zucchini-nlp  In addition to the above I found a couple of more things,

1) The padding_idx must be updated to None if the number of tokens in the new embeddings is smaller than the padding_idx. 

2) The model Lxmert has a bias term that has to be updated similar to the final_logits_bias in mbart. This if not updated would cause some failures in test cases. This problem has been there from before and I believe was not addressed.

In addition to what we already discussed about, I have made these changes in this commit as well.  All of the test cases are passing for me now.

@zucchini-nlp I am not sure what this tests_hub failure is. Could you please help out

@zucchini-nlp Made the changes that you have recommended. Please take a look. Thanks!

@bayllama thanks, everything looks good! Rerunning the tests didn't help, so can you merge `main` in case it was already fixed there

@amyeroberts can you merge this plz? Unrelated hub failures, from internal slack seems like we're not the only ones seeing it 

@zucchini-nlp There's been a fix push to main. @bayllama could you try rebasing? 

I've also re-requested review as there's been several commits since my approval

@zucchini-nlp @amyeroberts  Seems like the tests went through after rebasing. Let me know if anything else is required here.

@amyeroberts @zucchini-nlp Added the test to make sure the correct type is returned. Also made the other changes suggested.

This PR causes `model.resize_token_embeddings` to set `vocab_size` to zero (check this [thread](https://github.com/huggingface/transformers/issues/32170)). Reverting this PR resolved the issue. Could you help with this?

@zucchini-nlp @amyeroberts I believe zucchini-nlp already pushed in a fix for this. Let me know if I need to do something.