Skip to content

[UMT5] Ignore tied encoder embedding missing-key warning - #44014

Closed
weiguangli-io wants to merge 1 commit into
huggingface:mainfrom
weiguangli-io:codex/transformers-43992-umt5-embed-missing
Closed

[UMT5] Ignore tied encoder embedding missing-key warning#44014
weiguangli-io wants to merge 1 commit into
huggingface:mainfrom
weiguangli-io:codex/transformers-43992-umt5-embed-missing

Conversation

@weiguangli-io

Copy link
Copy Markdown
Contributor

Summary

Fixes #43992 by preventing a false missing-key report for UMT5EncoderModel when encoder.embed_tokens.weight is tied to shared.weight.

UMT5EncoderModel already declares tied weights, but loading checkpoints that only carry the shared embedding can still surface a misleading missing key warning for encoder.embed_tokens.weight. This change marks that key as ignorable-on-load for this encoder-only class, matching the tied-embedding semantics.

Changes

  • src/transformers/models/umt5/modeling_umt5.py
    • Add _keys_to_ignore_on_load_missing = [r"encoder.embed_tokens.weight"] to UMT5EncoderModel.
  • tests/models/umt5/test_encoder_missing_keys.py
    • Add a regression test that loads UMT5EncoderModel from a state dict missing encoder.embed_tokens.weight and asserts the key is not reported in missing_keys.

Validation

  • PYTHONPATH=src python3 -m ruff check src/transformers/models/umt5/modeling_umt5.py tests/models/umt5/test_encoder_missing_keys.py
  • PYTHONPATH=src /tmp/tfpycheck/bin/pytest -q tests/models/umt5/test_encoder_missing_keys.py
  • PYTHONPATH=src /tmp/tfpycheck/bin/coverage run -m pytest -q tests/models/umt5/test_encoder_missing_keys.py
  • PYTHONPATH=src /tmp/tfpycheck/bin/coverage report -m src/transformers/models/umt5/modeling_umt5.py

Coverage note:

  • The added regression test hits the new load-ignore behavior directly; full-file coverage remains low because the model file contains many unrelated execution paths.

@github-actions

Copy link
Copy Markdown
Contributor

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

run-slow: umt5

@Rocketknight1

Copy link
Copy Markdown
Member

I think this is a duplicate of #43880!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

UMT5Encoder.from_pretrained misses embed_tokens.weight

2 participants