Fixed loading LongT5 from legacy checkpoints
Author: Szustarol
Pull request: yes

# What does this PR do?

<!--
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 #40635

The problem with loading legacy LongT5 checkpoints stems from the way the shared embedding tokens are defined in the model: 
https://github.com/huggingface/transformers/blob/96a5774f2e048c21b9c0408c4fddfcf23a8934c5/src/transformers/models/longt5/modeling_longt5.py#L1350-L1356
In most legacy checkpoints they are not tied via the tied keys, but were still trained from the same shared weight that is created here:
https://github.com/huggingface/transformers/blob/96a5774f2e048c21b9c0408c4fddfcf23a8934c5/src/transformers/models/longt5/modeling_longt5.py#L1901

The solution to this problem is to manually assign the weights in overriden `from_pretrained` if the model contains the shared weights.
This fix is somewhat dirty, it would be much better to just update them directly in the state_dict to avoid the missing keys warning, but there is currently no public interface of the `PretrainedModel` class that allows for that.

If the reviewer has a better idea of how this could be accounted for I'm ready to devote my time towards fixing this issue.

## 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?
- [ ] 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.
- [x] 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?

@Rocketknight1 

<!-- 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, @qubvel
- speech models: @eustlb
- graph models: @clefourrier

Library:

- flax: @gante and @Rocketknight1
- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- pipelines: @Rocketknight1
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @zach-huggingface, @SunMarc and @qgallouedec
- chat templates: @Rocketknight1

Integrations:

- deepspeed: HF Trainer/Accelerate: @SunMarc @zach-huggingface
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc @MekkCyber

Documentation: @stevhliu

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: @Rocketknight1
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
Just a word of commentary that we might discuss as well: I believe `lm_head` (if it exists) should also be loaded from the shared weights looking at how the model would usually be trained, but this is not evident from the model's architecture. 

Hi @Szustarol it seems like you've identified the right issue, but is there a reason we can't just set `tied_weights_keys` in the modeling code?

Thanks for the follow-up @Rocketknight1, the `_tied_weights_keys` are actually set in the modelling code, but the weights are not saved under the name of any of the actually tied weights ("encoder.embed_tokens.weight", "decoder.embed_tokens.weight", or "lm_head.weight"). Instead, they are saved as "shared.weight", since this is the source parameter name.
We could add "shared.weight" to the _tied_weights_keys, which would nicely resolve the issue by tying all the layers that need to be tied, but then we would still have to _disrespect_ the config of the model. 
The weights are actually only tied if this line runs:
https://github.com/huggingface/transformers/blob/e481228f2f7d894981049dec263b6b096cafc57d/src/transformers/models/longt5/modeling_longt5.py#L1741-L1744
But most T5 checkpoints don't actually set this (i.e. `config.tie_word_embeddings=False`) for most models, I have also reviewed this for popular LongT5 models such as google/long-t5-local-base (see https://huggingface.co/google/long-t5-local-base/blob/main/config.json#L27). As a side note, this custom `_tie_weights` is actually not needed since modeling_utils tie decoder and encoder embeddings anyway if tie_word_embeddings=True, here:
https://github.com/huggingface/transformers/blob/896e9cea1ade521b2648f4798218550f6c72190c/src/transformers/modeling_utils.py#L2790

After reviewing the modeling code it seems to me like the weights are always shared through the "shared" layer, so I don't understand why someone made the choice to implement it that way, instead of just tying them.

Furthermore, the _tie_weights code seems to also be wrong when the lm_head should be tied too:
https://github.com/huggingface/transformers/blob/e481228f2f7d894981049dec263b6b096cafc57d/src/transformers/models/longt5/modeling_longt5.py#L1904-L1908
[...]
https://github.com/huggingface/transformers/blob/e481228f2f7d894981049dec263b6b096cafc57d/src/transformers/models/longt5/modeling_longt5.py#L1941-L1944
`lm_head` remains untied in the above example, even though it should be.

Seems like `_tie_weights` was added to multiple models as part of https://github.com/huggingface/transformers/commit/95020f208ed7c30895685af60ef3a791fb2d45ff, which made LongT5 nonfunctional.
To be honest, I'm not completely sure why, since this is done, as mentioned previously, in modeling_utils.


To retain maximum compatibility we could:
- remove the `_tie_weights` override which seems to just be be wrong
- and add `shared.weight` to the `_tied_weights_keys`
I will proceed with testing this solution, but this is interfering with the original source code, so I would appreciate some comment on this proposal. 

A similar issue might affect other models to which the `_tie_weights` override has been added.

@Szustarol updating the modeling code is totally okay, especially if part of it seems to be wrong. Models like LongT5 are not widely used and so it is very possible that a past update broke them and no-one realized.

I think if you make the update and you can show that:

1) Tests (including slow tests) for LongT5 still pass
2) Weight sharing now works correctly for LongT5 models on the Hub

that is enough evidence for us to accept the change. Either way, thank you for your careful approach to this PR!

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: longt5

Thank you for the clarification @Rocketknight1, I think I now have a more resilient solution.

After a bit more consideration and model inspection I noticed that some models are indeed saved erroneously (i.e. they don't contain the lm_head saved). This concerns the Google models, and only in the safetensors format (.bin works correctly).
Given that those models are the most used ones (judging by download numbers) I decided to implement a solution that takes care about all the possible use cases, but this means that the `lm_head` _cannot_ always be tied, because sometimes it is indeed supplied as a separate weight in the .bin or .safetensors checkpoint. 

The current solution involves iterating over missing keys and checking if they should be tied. If yes, a warning is emitted (as in this case, `tie_word_embeddings` should be `True`), and a fallback weight value is loaded from the `shared` weights.
This is the only solution that works for the three types of models on the hub:
 1) models with saved shared weights requiring `lm_head`, `decoder` and `encoder` tying,
 2) models with saved shared weights and `lm_head`, requirng `decoder` and `encoder` tying,
 3) models with saved `decoder`, `encoder` and optionally `lm_head` weights, requiring no tying.

Implementing the tying method would make the usecase 2) not functional, since 'lm_head' would be tied too. Also, the current solution allows us to keep the current `_tie_weights`, for whatever reason it was ever introduced.

I evaluated my changes on 4 models, with the following snippet:
```py
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
# model_id = "google/long-t5-tglobal-base"
model_id = "google/long-t5-local-base"
# model_id = "pszemraj/long-t5-tglobal-base-sci-simplify"
# model_id = "agemagician/mlong-t5-tglobal-base"
model = AutoModelForSeq2SeqLM.from_pretrained(
    model_id,
    use_safetensors=True
)

tokenizer = AutoTokenizer.from_pretrained(
    model_id,
)

text = "A brown fox did"
inputs = tokenizer(text, return_tensors="pt", max_length=20, truncation=True)
outputs = model.generate(**inputs, max_length=20)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
```

| model_id | Output - before the fix - .bin | Output - before the fix - .safetensor | Output - fixed - .bin | Output - fixed - .safetensor |
|---|---|---|---|---|
| google/long-t5-tglobal-base | Is it a fox or a fox terrier that is s | <empty response> | Is it a fox or a fox terrier that is s | Is it a fox or a fox terrier that is s |
|google/long-t5-local-base| ),),),),),),),),),),),),),),),),),),), | <empty response> | ),),),),),),),),),),),),),),),),),),), | ),),),),),),),),),),),),),),),),),),), |
| pszemraj/long-t5-tglobal-base-sci-simplify | Foxes were foxes that did not do their own things. | thea: and to of | Foxes were foxes that did not do their own things. | Foxes were foxes that did not do their own things. |
|agemagician/mlong-t5-tglobal-base|have a good time with the foxes. They are very good at laying eggs and|<empty response>| have a good time with the foxes. They are very good at laying eggs and|have a good time with the foxes. They are very good at laying eggs and|

Seems like "google/long-t5-tlocal-base" is corrupted, by the way.

Anyway, I think this wraps it up. 

PyTest report:
[report.xml](https://github.com/user-attachments/files/22221947/report.xml)


Seems good - cc @arthurzucker to confirm since this overrides `from_pretrained` for a text model

Dear reviewers (@Rocketknight1)
Since I see that the issue has become a bit stale, would you be more comfortable accepting a solution that overrides _load_pretrained_model instead? 
Initially, I considered overriding public methods more reasonable as their interface is less likely to have a breaking change, but of course _load_pretrained_model would work too, especially if you'd like the model to not report those legacy weights as missing after load.
Thanks a lot for any feedback!

@Szustarol no, it's okay! I think this solution is good - we can merge it, and someone can yell at me afterwards if it's a problem :sweat_smile: 

Thank you for the PR!

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