A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 21:45 UTC
Fix tied weights in weight mapping test for Transformers v5 (#36788)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Mar 11, 2026 at 15:10 UTC
d5816c8c2fa8dba84dc518c481a21bc6e5439acb
1 file changed
+9
-6
tests/models/multimodal/test_mapping.py
+9
-6
@@ -31,12 +31,6 @@ def create_dummy_model(repo: str, model_arch: str) -> PreTrainedModel:
31
config = AutoConfig.from_pretrained(repo)
32
with torch.device("meta"):
33
model = model_cls._from_config(config)
34
- # TODO(hmellor): Remove this once Transformers has fixed tied weights on meta device
35
- # https://github.com/huggingface/transformers/issues/43522
36
- if getattr(config.get_text_config(), "tie_word_embeddings", False) or getattr(
37
- config, "tie_word_embeddings", False
38
- ):
39
- model.tie_weights()
34
return model
35
36
@@ -103,6 +97,15 @@ def test_hf_model_weights_mapper(model_arch: str):
97
# Some checkpoints may have buffers, we ignore them for this test
98
ref_weight_names -= buffer_names
99
100
+ # Some checkpoints include tied weights (e.g. lm_head tied to embed_tokens) in the
101
+ # safetensors file. In Transformers v5, named_parameters() will not include them
102
+ # after they are tied in the model, so the mapper will not be able to map them.
103
+ # We exclude them from the reference weight names for this test.
104
+ if isinstance(tied := getattr(hf_dummy_model, "_tied_weights_keys", None), dict):
105
+ mapped_tied_weights = mapper.apply((k, None) for k in tied)
106
+ tied_weight_names = set(map(lambda x: x[0], mapped_tied_weights))
107
+ ref_weight_names -= tied_weight_names
108
+
109
weights_missing = ref_weight_names - weight_names
110
weights_unmapped = weight_names - ref_weight_names
111
assert not weights_missing and not weights_unmapped, (