[Bugfix][Model] Skip loading lm_head weights if using tie_word_embeddings

[tjohnson31415]

In llama and other models with tie_word_embeddings, there seems to be cases where the weight files will include both the lm_head.weight and embed_tokens.weight tensors, particularly after tuning procedures. Attempting to load such weights results in an error like:

[rank0]:   File "/opt/vllm/lib64/python3.11/site-packages/vllm/worker/model_runner.py", line 553, in load_model
[rank0]:     self.model = get_model(model_config=self.model_config,
[rank0]:                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/vllm/lib64/python3.11/site-packages/vllm/model_executor/model_loader/__init__.py", line 21, in get_model
[rank0]:     return loader.load_model(model_config=model_config,
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/vllm/lib64/python3.11/site-packages/vllm/model_executor/model_loader/loader.py", line 278, in load_model
[rank0]:     model.load_weights(
[rank0]:   File "/opt/vllm/lib64/python3.11/site-packages/vllm/model_executor/models/llama.py", line 481, in load_weights
[rank0]:     param = params_dict[name]
[rank0]:             ~~~~~~~~~~~^^^^^^
[rank0]: KeyError: 'lm_head.weight'

This error arises when the model does not have an lm_head.weight in named_parameters() but the weights files include lm_head.weight. With tie_word_embeddings set to true, lm_head.weight should be a duplicate of embed_tokens.weight, so the extra tensor seems to be included in the .safetensors unnecessarily. The change in this PR is to ignore lm_head.weight when tie_word_embeddings is true.

#3553 is a related issue that saw the same issue for Gemma which always uses tied weights.

I was made aware of this error in regards to a fine-tune of ibm-granite/granite-3b-code-instruct, which is a llama architecture with tie_word_embeddings set to True. After understanding the cause, I looked for other model implementations that may have the same issue with:

grep -r tie_word_embeddings vllm/model_executor/models/

I found some model implementations with tie_word_embeddings already include a check to skip loading lm_head.weight (qwen2, starcoder2, falcon). I added a check to the other models that did not include that check in load_weights.

I only tested the fix with llama, but included the checks for the models that were missing it to head off future issues. Let me know if it would be preferable to only change this for llama and leave the other models for a later PRs.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which consists a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of default ones by unblocking the steps in your fast-check build on Buildkite UI.

Once the PR is approved and ready to go, please make sure to run full CI as it is required to merge (or just use auto-merge).

To run full CI, you can do one of these:

• Comment /ready on the PR
• Add ready label to the PR
• Enable auto-merge.

🚀

[njhill]

Thanks @tjohnson31415!

[njhill]

@tjohnson31415 could you merge in latest main branch? I think that should fix the failing tests

[njhill]

@tjohnson31415 could you do it one more time? :) another fix went in for the tensorizer tests

[TissueC]

A emergent solution for this issue: manually initialize the lm_head with word_embedding, and set the tie_word_embeddings in config.json as False.
Hope this PR could be merged ASAP