"tests/vscode:/vscode.git/clone" did not exist on "602358f8a86ef9fc0ba882e083e19b44e00b9302"
Unverified Commit 14385c80 authored by Harry Mellor's avatar Harry Mellor Committed by GitHub
Browse files

Fix weight mapping test for Transfomers v5 (#33162)


Signed-off-by: default avatarHarry Mellor <19981378+hmellor@users.noreply.github.com>
parent 76139d08
......@@ -30,7 +30,12 @@ def create_dummy_model(repo: str, model_arch: str) -> PreTrainedModel:
model_cls: PreTrainedModel = getattr(transformers, model_arch)
config = AutoConfig.from_pretrained(repo)
with torch.device("meta"):
return model_cls._from_config(config)
model = model_cls._from_config(config)
# TODO(hmellor): Remove this once Transformers has fixed tied weights on meta device
# https://github.com/huggingface/transformers/issues/43522
if getattr(config.get_text_config(), "tie_word_embeddings", False):
model.tie_weights()
return model
def model_architectures_for_test() -> list[str]:
......
......@@ -249,7 +249,8 @@ class Base(
# Layers before module list
for name in pp_plan[:module_list_idx]:
if self.pp_group.is_first_rank or (
self.text_config.tie_word_embeddings and self.pp_group.is_last_rank
getattr(self.text_config, "tie_word_embeddings", False)
and self.pp_group.is_last_rank
):
continue
setattr(self.model, name, PPMissingLayer())
......
......@@ -38,7 +38,8 @@ class CausalMixin(VllmModelForTextGeneration):
# Tell `Base.load_weights` to skip
# `lm_head` if the model has tied word embeddings
if self.text_config.tie_word_embeddings:
tie_word_embeddings = getattr(self.text_config, "tie_word_embeddings", False)
if tie_word_embeddings:
self.skip_prefixes.append("lm_head.")
if self.pp_group.is_last_rank:
......@@ -48,7 +49,7 @@ class CausalMixin(VllmModelForTextGeneration):
quant_config=self.quant_config,
prefix=maybe_prefix(prefix, "lm_head"),
)
if self.text_config.tie_word_embeddings:
if tie_word_embeddings:
self.lm_head = self.lm_head.tie_weights(
self.model.get_input_embeddings()
)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment