From a37f07472fd4fd3272769a21d1bba718d8b2e5ca Mon Sep 17 00:00:00 2001 From: Will Date: Fri, 27 Sep 2024 19:02:24 -0400 Subject: [PATCH] fix: check for wte.weight along with embed_tokens.weight (#356) Signed-off-by: Will Johnson --- tuning/utils/merge_model_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tuning/utils/merge_model_utils.py b/tuning/utils/merge_model_utils.py index a2ec6577b..c2078908a 100644 --- a/tuning/utils/merge_model_utils.py +++ b/tuning/utils/merge_model_utils.py @@ -168,7 +168,7 @@ def post_process_vLLM_adapters_new_tokens( # vLLM requires renaming to output_embeddings new_embeddings["output_embeddings"] = new_output_embeddings - elif "embed_tokens.weight" in k: + elif "embed_tokens.weight" in k or "wte.weight" in k: embed_tokens = f.get_tensor(k) # pull out tensor values of new tokens new_input_embeddings = embed_tokens[-num_added_tokens:]