Fix UMT5EncoderModel embedding weights not being tied after loading

[jiqing-feng]

After the "Refactor weight loading" commit (#41580), UMT5EncoderModel fails to properly tie encoder.embed_tokens.weight to shared.weight when loading checkpoints with tie_word_embeddings=False (e.g., Wan-AI video generation models).

This causes encoder.embed_tokens.weight to remain uninitialized (all zeros), producing garbage output in downstream applications like WanPipeline.

Root Cause
The new tie_weights() in modeling_utils.py skips weight binding when tie_word_embeddings=False. However, UMT5EncoderModel architecturally requires encoder.embed_tokens to share weights with shared, regardless of the config setting.

Solution
Override tie_weights() in UMT5EncoderModel to always bind the embedding weights

Verification
Before fix:

import torch
from diffusers import AutoencoderKLWan, WanPipeline

model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)

text_encoder = pipe.text_encoder
print(f"text_encoder.shared.weight sum: {text_encoder.shared.weight.sum().item()}")
print(f"text_encoder.encoder.embed_tokens.weight sum: {text_encoder.encoder.embed_tokens.weight.sum().item()}")
print(f"Are they same tensor? {text_encoder.shared.weight.data_ptr() == text_encoder.encoder.embed_tokens.weight.data_ptr()}")

Before fix:
encoder.embed_tokens.weight.sum() = 0.0
shared.weight.sum() = -4456448.0
Same tensor: False

After fix:
Both weights sum to -4456448.0
Same tensor: True
WanPipeline video generation works correctly

The video script see huggingface/diffusers#13105

[zucchini-nlp]

This needs to be fixed in config file. We have hardcoded tie_word_embeddings=True tho if the config has tie_word_embeddings=False saved then it'll be in kwargs and override our hardcoded value

We need to pop tie_word_embeddings from kwargs if any

[zucchini-nlp]

same happens in other T5 family models, can you check them as well?

[jiqing-feng]

Hi @zucchini-nlp . Do you mean I should change the model config.json file like: https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B-Diffusers/discussions/15 ?

I am afraid I can't handle all t5 models cause there are so many models: https://huggingface.co/models?search=t5.

Please let me know if I misunderstood it. Thanks!

[zucchini-nlp]

No, no, you don't need to change the hub config. I meant the file config_t5.py needs to pop tie_word_embeddingd from kwarg to avoid overriding it in super call

[jiqing-feng]

Hi @zucchini-nlp . If I remove the current change of modeling and change the config file by self.tie_word_embeddings = kwargs.pop("tie_word_embeddings", True) , it will not work.
Do you mean I should keep the change in modeling file and also change the config file? (Just like my new change). Please let me know if I misunderstood this. Thanks!

[zucchini-nlp]

No, just pop without using like

kwargs.pop("tie_word_embeddings", None)
self.tie_word_embeddings = True # hardcode here for BC

[jiqing-feng]

Hi @zucchini-nlp . Thanks for your guide. I've changed it by your comment and verified it works.
For other T5 models, I am not sure if the tie weights should be changed like this. I will update it in a separate PR once I find bugs.

[zucchini-nlp]

Other T5 models also tie weights by default, and we have hardcoded True. So yeah, we can update all models

[jiqing-feng]

Hi @zucchini-nlp . I also updated other t5 models, please review it.

[zucchini-nlp]

let's revert Vilt and merge

[zucchini-nlp]

vilt doesn't force True, we need to revert this file

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[ArthurZucker]

I feel like we are missing a test no?

[zucchini-nlp]

Yep, a test would be nice as well!

[github-actions]

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

run-slow: mt5, umt5

[jiqing-feng]

Hi @zucchini-nlp . I have reverted the vilt change, and also added tests to make sure the tie_word_embeddings are forced to be True

[zucchini-nlp]

Thanks, let's merge!